File size: 5,434 Bytes
55accfa 8623ebb 55accfa 8623ebb 55accfa 8623ebb 55accfa 8623ebb 55accfa 8623ebb 55accfa 8623ebb 55accfa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import torch
import torch.nn as nn
import torch.nn.functional as F
def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> torch.Tensor:
grid_h = torch.arange(grid_size, dtype=torch.float32)
grid_w = torch.arange(grid_size, dtype=torch.float32)
grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
grid = torch.stack(grid, dim=0).reshape([2, 1, grid_size, grid_size])
dim = embed_dim // 2
omega = torch.arange(dim // 2, dtype=torch.float32) / (dim / 2.0)
omega = 1.0 / (10000 ** omega)
out_w = torch.einsum("m,d->md", grid[0].reshape(-1), omega)
out_h = torch.einsum("m,d->md", grid[1].reshape(-1), omega)
pos_embed = torch.cat([
torch.cat([torch.sin(out_w), torch.cos(out_w)], dim=1),
torch.cat([torch.sin(out_h), torch.cos(out_h)], dim=1)
], dim=1)
if cls_token:
pos_embed = torch.cat([torch.zeros([1, embed_dim], dtype=torch.float32), pos_embed], dim=0)
return pos_embed
class ConvNeXtStem(nn.Module):
"""Sub-pixel ConvNeXt Stem preserving high-frequency edge gradients."""
def __init__(self, in_channels: int = 3, embed_dim: int = 256):
super().__init__()
mid_dim = max(64, embed_dim // 2)
self.stem = nn.Sequential(
nn.Conv2d(in_channels, mid_dim, kernel_size=7, stride=2, padding=3),
nn.BatchNorm2d(mid_dim),
nn.GELU(),
nn.Conv2d(mid_dim, embed_dim, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(embed_dim),
nn.GELU(),
nn.Conv2d(embed_dim, embed_dim, kernel_size=4, stride=4),
nn.BatchNorm2d(embed_dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
feat = self.stem(x)
return feat.flatten(2).transpose(1, 2)
class TransformerBlock(nn.Module):
def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim, eps=1e-6)
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.qkv = nn.Linear(dim, dim * 3, bias=True)
self.proj = nn.Linear(dim, dim)
self.norm2 = nn.LayerNorm(dim, eps=1e-6)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = nn.Sequential(
nn.Linear(dim, mlp_hidden_dim),
nn.GELU(),
nn.Linear(mlp_hidden_dim, dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, N, C = x.shape
x_norm = self.norm1(x)
qkv = self.qkv(x_norm).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
attn_out = attn_out.transpose(1, 2).reshape(B, N, C)
x = x + self.proj(attn_out)
x = x + self.mlp(self.norm2(x))
return x
class AevanisVisionBackbone(nn.Module):
"""
Aevanis Vision Foundation Backbone (Universal Visual Eye)
- Nano: 7.68M parameters (embed_dim=256, depth=8, heads=4)
- Core: 24.35M parameters (embed_dim=384, depth=12, heads=6)
"""
def __init__(self, model_scale: str = "nano", img_size: int = 224, patch_size: int = 16, in_channels: int = 3):
super().__init__()
self.model_scale = model_scale
if model_scale == "nano":
embed_dim, depth, num_heads = 256, 8, 4
elif model_scale == "core":
embed_dim, depth, num_heads = 384, 12, 6
else:
raise ValueError(f"Unknown scale: {model_scale}")
self.embed_dim = embed_dim
self.patch_embed = ConvNeXtStem(in_channels=in_channels, embed_dim=embed_dim)
grid_size = img_size // patch_size
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
pos_embed = get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=True)
self.register_buffer("pos_embed", pos_embed.unsqueeze(0))
self.blocks = nn.ModuleList([
TransformerBlock(embed_dim, num_heads, mlp_ratio=4.0) for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim, eps=1e-6)
def forward(self, x: torch.Tensor):
B = x.shape[0]
x = self.patch_embed(x)
x = x + self.pos_embed[:, 1:, :]
cls_tokens = self.cls_token.expand(B, -1, -1) + self.pos_embed[:, :1, :]
x = torch.cat((cls_tokens, x), dim=1)
for blk in self.blocks:
x = blk(x)
x = self.norm(x)
cls_token = x[:, 0]
spatial_tokens = x[:, 1:] # [B, 196, embed_dim]
return {
"cls_token": cls_token,
"spatial_tokens": spatial_tokens,
"full_tokens": x
}
def forward_encoder(self, x: torch.Tensor, mask_ratio: float = 0.0):
# Backward-compatibility alias
out = self.forward(x)
return out["full_tokens"], None, None
@classmethod
def from_pretrained(cls, weight_path: str, model_scale: str = "nano"):
model = cls(model_scale=model_scale)
ckpt = torch.load(weight_path, map_location="cpu")
cleaned = {k.replace("module.", ""): v for k, v in ckpt.items() if not k.startswith("decoder")}
model.load_state_dict(cleaned, strict=False)
model.eval()
return model
|