Spaces:
Running
on
L40S
Running
on
L40S
| import math | |
| import torch | |
| import torch.nn as nn | |
| from cube3d.model.transformers.norm import LayerNorm, RMSNorm | |
| def init_linear(module, embed_dim: int): | |
| """ | |
| Initializes the weights and biases of a given linear module. | |
| Args: | |
| module (nn.Module): The module to initialize. Expected to be an instance of nn.Linear. | |
| embed_dim (int): The embedding dimension used to calculate the standard deviation | |
| for weight initialization. | |
| Returns: | |
| None | |
| """ | |
| if isinstance(module, nn.Linear): | |
| nn.init.normal_(module.weight, std=math.sqrt(1.0 / embed_dim)) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| def init_tfixup(module: nn.Module, num_layers: int): | |
| """Special initialization from https://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf | |
| Args: | |
| module (nn.Module): decoder/encoder module | |
| num_layers (int): number of layers in the module | |
| """ | |
| with torch.no_grad(): | |
| for pn, p in module.named_parameters(): | |
| if ( | |
| pn.endswith("c_proj.weight") | |
| or pn.endswith("up_proj.weight") | |
| or pn.endswith("down_proj.weight") | |
| ): | |
| p *= (4 * num_layers) ** (-0.25) | |
| elif pn.endswith("c_v.weight"): | |
| p *= (4 * num_layers) ** (-0.25) * math.sqrt(2) | |
| class MLP(nn.Module): | |
| def __init__(self, embed_dim, hidden_dim, bias=True, approximate="none"): | |
| """ | |
| MLP with GELU activation function." | |
| """ | |
| super().__init__() | |
| self.up_proj = nn.Linear(embed_dim, hidden_dim, bias=bias) | |
| self.down_proj = nn.Linear(hidden_dim, embed_dim, bias=bias) | |
| self.act_fn = nn.GELU(approximate=approximate) | |
| def forward(self, x): | |
| return self.down_proj(self.act_fn(self.up_proj(x))) | |
| class SelfAttention(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim: int, | |
| num_heads: int, | |
| bias: bool = True, | |
| eps: float = 1e-6, | |
| ): | |
| """ | |
| Initializes the self attention mechanism. | |
| Args: | |
| embed_dim (int): The dimensionality of the embedding space. | |
| num_heads (int): The number of attention heads. | |
| bias (bool, optional): Whether to include bias terms in the linear layers. Defaults to True. | |
| eps (float, optional): A small value added for numerical stability. Defaults to 1e-6. | |
| Raises: | |
| AssertionError: If `embed_dim` is not divisible by `num_heads`. | |
| """ | |
| super().__init__() | |
| assert embed_dim % num_heads == 0 | |
| self.num_heads = num_heads | |
| self.c_qk = nn.Linear(embed_dim, 2 * embed_dim, bias=bias) | |
| self.c_v = nn.Linear(embed_dim, embed_dim, bias=bias) | |
| self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias) | |
| head_dim = embed_dim // num_heads | |
| self.q_norm = RMSNorm(head_dim) | |
| self.k_norm = RMSNorm(head_dim) | |
| def forward(self, x, attn_mask=None, is_causal: bool = False): | |
| """ | |
| Performs the forward pass of the attention mechanism. | |
| Args: | |
| x (torch.Tensor): Input tensor. | |
| attn_mask (Optional[torch.Tensor]): Attention mask to apply. Default is None. | |
| is_causal (bool): If True, applies a causal mask to prevent attending to future positions. | |
| Default is False. | |
| Returns: | |
| torch.Tensor: Output tensor after applying | |
| the attention mechanism and projection. | |
| """ | |
| b, l, d = x.shape | |
| q, k = self.c_qk(x).chunk(2, dim=-1) | |
| v = self.c_v(x) | |
| q = q.view(b, l, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| k = k.view(b, l, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| v = v.view(b, l, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| is_causal = is_causal and attn_mask is None | |
| y = torch.nn.functional.scaled_dot_product_attention( | |
| q, | |
| k, | |
| v, | |
| attn_mask=attn_mask, | |
| dropout_p=0.0, | |
| is_causal=is_causal, | |
| ) | |
| y = y.transpose(1, 2).contiguous().view(b, l, d) | |
| y = self.c_proj(y) | |
| return y | |
| class CrossAttention(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim: int, | |
| num_heads: int, | |
| q_dim=None, | |
| kv_dim=None, | |
| bias: bool = True, | |
| ): | |
| """ | |
| Initializes the cross attention mechanism. | |
| Args: | |
| embed_dim (int): The dimensionality of the embedding space. | |
| num_heads (int): The number of attention heads. | |
| q_dim (int, optional): The dimensionality of the query input. Defaults to `embed_dim`. | |
| kv_dim (int, optional): The dimensionality of the key and value inputs. Defaults to `embed_dim`. | |
| bias (bool, optional): Whether to include a bias term in the linear projections. Defaults to True. | |
| Raises: | |
| AssertionError: If `embed_dim` is not divisible by `num_heads`. | |
| """ | |
| super().__init__() | |
| assert embed_dim % num_heads == 0 | |
| q_dim = q_dim or embed_dim | |
| kv_dim = kv_dim or embed_dim | |
| self.c_q = nn.Linear(q_dim, embed_dim, bias=bias) | |
| self.c_k = nn.Linear(kv_dim, embed_dim, bias=bias) | |
| self.c_v = nn.Linear(kv_dim, embed_dim, bias=bias) | |
| self.c_proj = nn.Linear(embed_dim, embed_dim, bias=bias) | |
| self.num_heads = num_heads | |
| def forward(self, x, c, attn_mask=None, is_causal: bool = False): | |
| """ | |
| Forward pass for the attention mechanism. | |
| Args: | |
| x (torch.Tensor): Input tensor of shape. | |
| c (torch.Tensor): Context tensor. | |
| attn_mask (torch.Tensor, optional): Attention mask. | |
| Defaults to None. | |
| is_causal (bool, optional): Whether to apply causal masking. Defaults to False. | |
| Returns: | |
| torch.Tensor: Output tensor. | |
| """ | |
| q, k = self.c_q(x), self.c_k(c) | |
| v = self.c_v(c) | |
| b, l, d = q.shape | |
| s = k.shape[1] | |
| q = q.view(b, l, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| k = k.view(b, s, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| v = v.view(b, s, self.num_heads, -1).transpose(1, 2) # (B, nh, T, hs) | |
| y = torch.nn.functional.scaled_dot_product_attention( | |
| q, | |
| k, | |
| v, | |
| attn_mask=attn_mask, | |
| dropout_p=0.0, | |
| is_causal=(attn_mask is not None) and is_causal, | |
| ) | |
| y = y.transpose(1, 2).contiguous().view(b, l, d) | |
| y = self.c_proj(y) | |
| return y | |
| class EncoderLayer(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim: int, | |
| num_heads: int, | |
| bias: bool = True, | |
| eps: float = 1e-6, | |
| ) -> None: | |
| """ | |
| Initializes the EncoderLayer module. | |
| Args: | |
| embed_dim (int): The dimensionality of the embedding space. | |
| num_heads (int): The number of attention heads. | |
| bias (bool, optional): Whether to include bias terms in the layers. Defaults to True. | |
| eps (float, optional): A small value added for numerical stability in normalization layers. Defaults to 1e-6. | |
| """ | |
| super().__init__() | |
| self.ln_1 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps) | |
| self.attn = SelfAttention(embed_dim, num_heads, bias=bias, eps=eps) | |
| self.ln_2 = LayerNorm(embed_dim, elementwise_affine=False, eps=eps) | |
| self.mlp = MLP(embed_dim=embed_dim, hidden_dim=embed_dim * 4, bias=bias) | |
| def forward(self, x, attn_mask=None, is_causal: bool = False): | |
| """ | |
| Performs the forward pass of the transformer block. | |
| Args: | |
| x (torch.Tensor): The input tensor. | |
| attn_mask (torch.Tensor, optional): An optional attention mask tensor to apply during the | |
| attention computation. Default is None. | |
| is_causal (bool, optional): If True, applies a causal mask to prevent attention to future | |
| positions. Default is False. | |
| Returns: | |
| torch.Tensor: The output tensor of the same shape as the input. | |
| """ | |
| x = x + self.attn(self.ln_1(x), attn_mask=attn_mask, is_causal=is_causal) | |
| x = x + self.mlp(self.ln_2(x)) | |
| return x | |
| class EncoderCrossAttentionLayer(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim: int, | |
| num_heads: int, | |
| q_dim=None, | |
| kv_dim=None, | |
| bias: bool = True, | |
| eps: float = 1e-6, | |
| ) -> None: | |
| """ | |
| Initializes the EncoderAttentionLayer module with cross-attention, | |
| and a feed-forward MLP. | |
| Args: | |
| embed_dim (int): The dimensionality of the embedding space. | |
| num_heads (int): The number of attention heads. | |
| q_dim (int, optional): Dimensionality of the query input. Defaults to `embed_dim`. | |
| kv_dim (int, optional): Dimensionality of the key and value inputs. Defaults to `embed_dim`. | |
| bias (bool, optional): Whether to include bias terms in the layers. Defaults to True. | |
| eps (float, optional): A small value added to the denominator for numerical stability | |
| in layer normalization. Defaults to 1e-6. | |
| """ | |
| super().__init__() | |
| q_dim = q_dim or embed_dim | |
| kv_dim = kv_dim or embed_dim | |
| self.attn = CrossAttention( | |
| embed_dim, | |
| num_heads, | |
| q_dim=q_dim, | |
| kv_dim=kv_dim, | |
| bias=bias, | |
| ) | |
| self.ln_1 = LayerNorm(q_dim, elementwise_affine=False, eps=eps) | |
| self.ln_2 = LayerNorm(kv_dim, elementwise_affine=False, eps=eps) | |
| self.ln_f = LayerNorm(embed_dim, elementwise_affine=False, eps=eps) | |
| self.mlp = MLP(embed_dim=embed_dim, hidden_dim=embed_dim * 4, bias=bias) | |
| def forward(self, x, c, attn_mask=None, is_causal: bool = False): | |
| """ | |
| Forward pass for the attention mechanism. | |
| Args: | |
| x (torch.Tensor): The input tensor to the attention mechanism. | |
| c (torch.Tensor): The context tensor used for cross-attention. | |
| attn_mask (torch.Tensor, optional): An optional attention mask to control | |
| which positions can attend to others. Defaults to None. | |
| is_causal (bool, optional): If True, applies a causal mask to prevent | |
| attending to future positions. Defaults to False. | |
| Returns: | |
| torch.Tensor: The output tensor after applying attention and MLP layers. | |
| """ | |
| x = x + self.attn( | |
| self.ln_1(x), self.ln_2(c), attn_mask=attn_mask, is_causal=is_causal | |
| ) | |
| x = x + self.mlp(self.ln_f(x)) | |
| return x | |