|
from typing import *
|
|
import torch
|
|
import torch.nn as nn
|
|
from ...modules.utils import convert_module_to_f16, convert_module_to_f32
|
|
from ...modules import sparse as sp
|
|
from ...modules.transformer import AbsolutePositionEmbedder
|
|
from ...modules.sparse.transformer import SparseTransformerBlock
|
|
|
|
|
|
def block_attn_config(self):
|
|
"""
|
|
Return the attention configuration of the model.
|
|
"""
|
|
for i in range(self.num_blocks):
|
|
if self.attn_mode == "shift_window":
|
|
yield "serialized", self.window_size, 0, (16 * (i % 2),) * 3, sp.SerializeMode.Z_ORDER
|
|
elif self.attn_mode == "shift_sequence":
|
|
yield "serialized", self.window_size, self.window_size // 2 * (i % 2), (0, 0, 0), sp.SerializeMode.Z_ORDER
|
|
elif self.attn_mode == "shift_order":
|
|
yield "serialized", self.window_size, 0, (0, 0, 0), sp.SerializeModes[i % 4]
|
|
elif self.attn_mode == "full":
|
|
yield "full", None, None, None, None
|
|
elif self.attn_mode == "swin":
|
|
yield "windowed", self.window_size, None, self.window_size // 2 * (i % 2), None
|
|
|
|
|
|
class SparseTransformerBase(nn.Module):
|
|
"""
|
|
Sparse Transformer without output layers.
|
|
Serve as the base class for encoder and decoder.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
model_channels: int,
|
|
num_blocks: int,
|
|
num_heads: Optional[int] = None,
|
|
num_head_channels: Optional[int] = 64,
|
|
mlp_ratio: float = 4.0,
|
|
attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
|
|
window_size: Optional[int] = None,
|
|
pe_mode: Literal["ape", "rope"] = "ape",
|
|
use_fp16: bool = False,
|
|
use_checkpoint: bool = False,
|
|
qk_rms_norm: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.model_channels = model_channels
|
|
self.num_blocks = num_blocks
|
|
self.window_size = window_size
|
|
self.num_heads = num_heads or model_channels // num_head_channels
|
|
self.mlp_ratio = mlp_ratio
|
|
self.attn_mode = attn_mode
|
|
self.pe_mode = pe_mode
|
|
self.use_fp16 = use_fp16
|
|
self.use_checkpoint = use_checkpoint
|
|
self.qk_rms_norm = qk_rms_norm
|
|
self.dtype = torch.float16 if use_fp16 else torch.float32
|
|
|
|
if pe_mode == "ape":
|
|
self.pos_embedder = AbsolutePositionEmbedder(model_channels)
|
|
|
|
self.input_layer = sp.SparseLinear(in_channels, model_channels)
|
|
self.blocks = nn.ModuleList([
|
|
SparseTransformerBlock(
|
|
model_channels,
|
|
num_heads=self.num_heads,
|
|
mlp_ratio=self.mlp_ratio,
|
|
attn_mode=attn_mode,
|
|
window_size=window_size,
|
|
shift_sequence=shift_sequence,
|
|
shift_window=shift_window,
|
|
serialize_mode=serialize_mode,
|
|
use_checkpoint=self.use_checkpoint,
|
|
use_rope=(pe_mode == "rope"),
|
|
qk_rms_norm=self.qk_rms_norm,
|
|
)
|
|
for attn_mode, window_size, shift_sequence, shift_window, serialize_mode in block_attn_config(self)
|
|
])
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
"""
|
|
Return the device of the model.
|
|
"""
|
|
return next(self.parameters()).device
|
|
|
|
def convert_to_fp16(self) -> None:
|
|
"""
|
|
Convert the torso of the model to float16.
|
|
"""
|
|
self.blocks.apply(convert_module_to_f16)
|
|
|
|
def convert_to_fp32(self) -> None:
|
|
"""
|
|
Convert the torso of the model to float32.
|
|
"""
|
|
self.blocks.apply(convert_module_to_f32)
|
|
|
|
def initialize_weights(self) -> None:
|
|
|
|
def _basic_init(module):
|
|
if isinstance(module, nn.Linear):
|
|
torch.nn.init.xavier_uniform_(module.weight)
|
|
if module.bias is not None:
|
|
nn.init.constant_(module.bias, 0)
|
|
self.apply(_basic_init)
|
|
|
|
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
|
h = self.input_layer(x)
|
|
if self.pe_mode == "ape":
|
|
h = h + self.pos_embedder(x.coords[:, 1:])
|
|
h = h.type(self.dtype)
|
|
for block in self.blocks:
|
|
h = block(h)
|
|
return h
|
|
|
