aimv2-4b-ve / modeling_aimv2.py
lucasjin's picture
Upload folder using huggingface_hub
08b3cf0 verified
from typing import Optional, Tuple, Union
import torch
from .configuration_aimv2 import AIMv2Config
from torch import nn
from torch.nn import functional as F
from transformers.modeling_outputs import BaseModelOutputWithNoAttention
from transformers.modeling_utils import PreTrainedModel
__all__ = ["AIMv2Model"]
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
output = self._norm(x.float()).type_as(x)
return output * self.weight
def extra_repr(self) -> str:
return f"{tuple(self.weight.shape)}, eps={self.eps}"
def _norm(self, x: torch.Tensor) -> torch.Tensor:
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
class AIMv2SwiGLUFFN(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
hidden_features = config.intermediate_size
in_features = config.hidden_size
bias = config.use_bias
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
self.fc2 = nn.Linear(hidden_features, in_features, bias=bias)
self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = F.silu(self.fc1(x)) * self.fc3(x)
x = self.fc2(x)
return x
class AIMv2PatchEmbed(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
self.proj = nn.Conv2d(
config.num_channels,
config.hidden_size,
kernel_size=(config.patch_size, config.patch_size),
stride=(config.patch_size, config.patch_size),
)
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.proj(x).flatten(2).transpose(1, 2)
x = self.norm(x)
return x
class AIMv2ViTPreprocessor(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
num_patches = (config.image_size // config.patch_size) ** 2
self.patchifier = AIMv2PatchEmbed(config)
self.pos_embed = nn.Parameter(torch.zeros((1, num_patches, config.hidden_size)))
def forward(self, x: torch.Tensor) -> torch.Tensor:
tokens = self.patchifier(x)
_, N, _ = tokens.shape
pos_embed = self.pos_embed.to(tokens.device)
tokens = tokens + pos_embed[:, :N]
return tokens
class AIMv2Attention(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
dim = config.hidden_size
self.num_heads = config.num_attention_heads
self.qkv = nn.Linear(dim, dim * 3, bias=config.qkv_bias)
self.attn_drop = nn.Dropout(config.attention_dropout)
self.proj = nn.Linear(dim, dim, bias=config.use_bias)
self.proj_drop = nn.Dropout(config.projection_dropout)
def forward(
self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
B, N, C = x.shape
qkv = (
self.qkv(x)
.reshape(B, N, 3, self.num_heads, C // self.num_heads)
.permute(2, 0, 3, 1, 4)
)
q, k, v = qkv.unbind(0)
x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
x = x.transpose(1, 2).contiguous().reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class AIMv2Block(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
self.attn = AIMv2Attention(config)
self.norm_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp = AIMv2SwiGLUFFN(config)
self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
x = x + self.attn(self.norm_1(x), mask)
x = x + self.mlp(self.norm_2(x))
return x
class AIMv2Transformer(nn.Module):
def __init__(self, config: AIMv2Config):
super().__init__()
self.blocks = nn.ModuleList(
[AIMv2Block(config) for _ in range(config.num_hidden_layers)]
)
self.post_trunk_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self,
tokens: torch.Tensor,
mask: Optional[torch.Tensor] = None,
output_hidden_states: bool = False,
) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]]]:
hidden_states = () if output_hidden_states else None
for block in self.blocks:
tokens = block(tokens, mask)
if output_hidden_states:
hidden_states += (tokens,)
tokens = self.post_trunk_norm(tokens)
return tokens, hidden_states
class AIMv2PretrainedModel(PreTrainedModel):
config_class = AIMv2Config
base_model_prefix = "aimv2"
main_input_name = "pixel_values"
_supports_sdpa = True
class AIMv2Model(AIMv2PretrainedModel):
def __init__(self, config: AIMv2Config):
super().__init__(config)
self.preprocessor = AIMv2ViTPreprocessor(config)
self.trunk = AIMv2Transformer(config)
def forward(
self,
pixel_values: torch.Tensor,
mask: Optional[torch.Tensor] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[
Tuple[torch.Tensor],
Tuple[torch.Tensor, Tuple[torch.Tensor, ...]],
BaseModelOutputWithNoAttention,
]:
if output_hidden_states is None:
output_hidden_states = self.config.output_hidden_states
if return_dict is None:
return_dict = self.config.use_return_dict
x = self.preprocessor(pixel_values)
x, hidden_states = self.trunk(
x, mask, output_hidden_states=output_hidden_states
)
if not return_dict:
res = (x,)
res += (hidden_states,) if output_hidden_states else ()
return res
return BaseModelOutputWithNoAttention(
last_hidden_state=x,
hidden_states=hidden_states,
)
from functools import partial
from torch import nn
import torch.nn.functional as F
from transformers.activations import ACT2FN
import math
import torch
class GLU(nn.Module):
def __init__(self, hidden_size, ffn_hidden_size, in_features):
super().__init__()
self.linear_proj = nn.Linear(in_features, hidden_size, bias=False)
self.norm1 = nn.LayerNorm(hidden_size)
self.act1 = nn.GELU()
self.act2 = nn.functional.silu
self.dense_h_to_4h = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
self.gate_proj = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
self.dense_4h_to_h = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
def forward(self, x):
x = self.linear_proj(x)
x = self.act1(self.norm1(x))
x = self.act2(self.gate_proj(x)) * self.dense_h_to_4h(x)
x = self.dense_4h_to_h(x)
return x
class MlpGLU(nn.Module):
def __init__(self, in_hidden_size, out_hidden_size):
super(MlpGLU, self).__init__()
ffn_hidden_size = out_hidden_size * 4 # out_hidden_size * 4 3584 * 4 = 14336
self.linear_proj = GLU(
hidden_size=out_hidden_size,
ffn_hidden_size=ffn_hidden_size,
in_features=in_hidden_size,
)
def forward(self, x, attention_mask: torch.Tensor = None):
x = self.linear_proj(x)
return x
class PixelShuffleLayer(nn.Module):
def __init(self):
super(PixelShuffleLayer, self).__init__()
def forward(self, x, scale_factor=0.5):
# print(f'in pixelshuffle: {x.shape}')
n, w, h, c = x.size()
# N, W, H, C --> N, W, H * scale, C // scale
x = x.reshape(n, w, int(h * scale_factor), int(c / scale_factor))
# N, W, H * scale, C // scale --> N, H * scale, W, C // scale
x = x.permute(0, 2, 1, 3).contiguous()
# N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
x = x.view(n, int(h * scale_factor), int(w * scale_factor),
int(c / (scale_factor * scale_factor)))
x = x.permute(0, 2, 1, 3).contiguous()
return x
class PixelShuffleConnector(nn.Module):
def __init__(self, in_hidden_size, out_hidden_size, down_rate=2):
super(PixelShuffleConnector, self).__init__()
# ffn_hidden_size = 13696
ffn_hidden_size = out_hidden_size * 4 # out_hidden_size * 4 3584 * 4 = 14336
self.linear_proj = GLU(
hidden_size=out_hidden_size,
ffn_hidden_size=ffn_hidden_size,
in_features=in_hidden_size * 4,
)
self.down_rate = down_rate
if self.down_rate == 2:
down = PixelShuffleLayer()
self.downsample = nn.Sequential(*[down])
else:
print(f"unsupported downsample rate for now!")
self.scaling_factor = 8
def forward(self, x, attention_mask: torch.Tensor = None):
# print(f'xin: {x.shape}')
b, s, h = x.shape
grid_size = int(s**0.5)
x = x.reshape(b, grid_size, grid_size, h)
x = self.downsample(x)
# print(f'x: {x.shape}')
# [11, 16, 16, 4608]
x = x.reshape(x.shape[0], -1, x.shape[-1])
x = self.linear_proj(x)
return x