| """SHViT architecture — Single-Head Vision Transformer. |
| |
| Source: https://github.com/novendrastywn/MicroViT (model/shvit.py) |
| Included verbatim; do not edit — keep in sync with upstream if needed. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import itertools |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from timm.models.vision_transformer import trunc_normal_ |
| from timm.models.layers import SqueezeExcite |
|
|
|
|
| class GroupNorm(nn.GroupNorm): |
| def __init__(self, num_channels, **kwargs): |
| super().__init__(1, num_channels, **kwargs) |
|
|
|
|
| class Conv2d_BN(nn.Sequential): |
| def __init__(self, a, b, ks=1, stride=1, pad=0, dilation=1, |
| groups=1, bn_weight_init=1): |
| super().__init__() |
| self.add_module('c', nn.Conv2d(a, b, ks, stride, pad, dilation, groups, bias=False)) |
| self.add_module('bn', nn.BatchNorm2d(b)) |
| nn.init.constant_(self.bn.weight, bn_weight_init) |
| nn.init.constant_(self.bn.bias, 0) |
|
|
| @torch.no_grad() |
| def fuse(self): |
| c, bn = self._modules.values() |
| w = bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| w = c.weight * w[:, None, None, None] |
| b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| m = nn.Conv2d( |
| w.size(1) * self.c.groups, w.size(0), w.shape[2:], |
| stride=self.c.stride, padding=self.c.padding, |
| dilation=self.c.dilation, groups=self.c.groups, |
| device=c.weight.device, |
| ) |
| m.weight.data.copy_(w) |
| m.bias.data.copy_(b) |
| return m |
|
|
|
|
| class BN_Linear(nn.Sequential): |
| def __init__(self, a, b, bias=True, std=0.02): |
| super().__init__() |
| self.add_module('bn', nn.BatchNorm1d(a)) |
| self.add_module('l', nn.Linear(a, b, bias=bias)) |
| trunc_normal_(self.l.weight, std=std) |
| if bias: |
| nn.init.constant_(self.l.bias, 0) |
|
|
| @torch.no_grad() |
| def fuse(self): |
| bn, l = self._modules.values() |
| w = bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| b = bn.bias - self.bn.running_mean * self.bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| w = l.weight * w[None, :] |
| if l.bias is None: |
| b = b @ self.l.weight.T |
| else: |
| b = (l.weight @ b[:, None]).view(-1) + self.l.bias |
| m = nn.Linear(w.size(1), w.size(0)) |
| m.weight.data.copy_(w) |
| m.bias.data.copy_(b) |
| return m |
|
|
|
|
| class PatchMerging(nn.Module): |
| def __init__(self, dim, out_dim): |
| super().__init__() |
| hid_dim = int(dim * 4) |
| self.conv1 = Conv2d_BN(dim, hid_dim, 1, 1, 0) |
| self.act = nn.ReLU() |
| self.conv2 = Conv2d_BN(hid_dim, hid_dim, 3, 2, 1, groups=hid_dim) |
| self.se = SqueezeExcite(hid_dim, 0.25) |
| self.conv3 = Conv2d_BN(hid_dim, out_dim, 1, 1, 0) |
|
|
| def forward(self, x): |
| return self.conv3(self.se(self.act(self.conv2(self.act(self.conv1(x)))))) |
|
|
|
|
| class Residual(nn.Module): |
| def __init__(self, m, drop=0.0): |
| super().__init__() |
| self.m = m |
| self.drop = drop |
|
|
| def forward(self, x): |
| if self.training and self.drop > 0: |
| return x + self.m(x) * torch.rand( |
| x.size(0), 1, 1, 1, device=x.device |
| ).ge_(self.drop).div(1 - self.drop).detach() |
| return x + self.m(x) |
|
|
| @torch.no_grad() |
| def fuse(self): |
| if isinstance(self.m, Conv2d_BN): |
| m = self.m.fuse() |
| assert m.groups == m.in_channels |
| identity = torch.ones(m.weight.shape[0], m.weight.shape[1], 1, 1) |
| identity = nn.functional.pad(identity, [1, 1, 1, 1]) |
| m.weight += identity.to(m.weight.device) |
| return m |
| return self |
|
|
|
|
| class FFN(nn.Module): |
| def __init__(self, ed, h): |
| super().__init__() |
| self.pw1 = Conv2d_BN(ed, h) |
| self.act = nn.ReLU() |
| self.pw2 = Conv2d_BN(h, ed, bn_weight_init=0) |
|
|
| def forward(self, x): |
| return self.pw2(self.act(self.pw1(x))) |
|
|
|
|
| class SHSA(nn.Module): |
| """Single-Head Self-Attention.""" |
|
|
| def __init__(self, dim, qk_dim, pdim): |
| super().__init__() |
| self.scale = qk_dim ** -0.5 |
| self.qk_dim = qk_dim |
| self.dim = dim |
| self.pdim = pdim |
| self.pre_norm = GroupNorm(pdim) |
| self.qkv = Conv2d_BN(pdim, qk_dim * 2 + pdim) |
| self.proj = nn.Sequential(nn.ReLU(), Conv2d_BN(dim, dim, bn_weight_init=0)) |
|
|
| def forward(self, x): |
| B, C, H, W = x.shape |
| x1, x2 = torch.split(x, [self.pdim, self.dim - self.pdim], dim=1) |
| x1 = self.pre_norm(x1) |
| qkv = self.qkv(x1) |
| q, k, v = qkv.split([self.qk_dim, self.qk_dim, self.pdim], dim=1) |
| q, k, v = q.flatten(2), k.flatten(2), v.flatten(2) |
| attn = (q.transpose(-2, -1) @ k) * self.scale |
| attn = attn.softmax(dim=-1) |
| x1 = (v @ attn.transpose(-2, -1)).reshape(B, self.pdim, H, W) |
| return self.proj(torch.cat([x1, x2], dim=1)) |
|
|
|
|
| class BasicBlock(nn.Module): |
| def __init__(self, dim, qk_dim, pdim, type): |
| super().__init__() |
| if type == "s": |
| self.conv = Residual(Conv2d_BN(dim, dim, 3, 1, 1, groups=dim, bn_weight_init=0)) |
| self.mixer = Residual(SHSA(dim, qk_dim, pdim)) |
| self.ffn = Residual(FFN(dim, int(dim * 2))) |
| elif type == "i": |
| self.conv = Residual(Conv2d_BN(dim, dim, 3, 1, 1, groups=dim, bn_weight_init=0)) |
| self.mixer = nn.Identity() |
| self.ffn = Residual(FFN(dim, int(dim * 2))) |
|
|
| def forward(self, x): |
| return self.ffn(self.mixer(self.conv(x))) |
|
|
|
|
| class SHViT(nn.Module): |
| def __init__( |
| self, |
| in_chans: int = 3, |
| num_classes: int = 1000, |
| embed_dim: list[int] = [128, 256, 384], |
| partial_dim: list[int] = [32, 64, 96], |
| qk_dim: list[int] = [16, 16, 16], |
| depth: list[int] = [1, 2, 3], |
| types: list[str] = ["s", "s", "s"], |
| down_ops: list = [["subsample", 2], ["subsample", 2], [""]], |
| distillation: bool = False, |
| ): |
| super().__init__() |
|
|
| self.patch_embed = nn.Sequential( |
| Conv2d_BN(in_chans, embed_dim[0] // 8, 3, 2, 1), nn.ReLU(), |
| Conv2d_BN(embed_dim[0] // 8, embed_dim[0] // 4, 3, 2, 1), nn.ReLU(), |
| Conv2d_BN(embed_dim[0] // 4, embed_dim[0] // 2, 3, 2, 1), nn.ReLU(), |
| Conv2d_BN(embed_dim[0] // 2, embed_dim[0], 3, 2, 1), |
| ) |
|
|
| self.blocks1: list = [] |
| self.blocks2: list = [] |
| self.blocks3: list = [] |
|
|
| for i, (ed, kd, pd, dpth, do, t) in enumerate( |
| zip(embed_dim, qk_dim, partial_dim, depth, down_ops, types) |
| ): |
| for _ in range(dpth): |
| getattr(self, f"blocks{i + 1}").append(BasicBlock(ed, kd, pd, t)) |
| if do[0] == "subsample": |
| blk = getattr(self, f"blocks{i + 2}") |
| blk.append(nn.Sequential( |
| Residual(Conv2d_BN(embed_dim[i], embed_dim[i], 3, 1, 1, groups=embed_dim[i])), |
| Residual(FFN(embed_dim[i], int(embed_dim[i] * 2))), |
| )) |
| blk.append(PatchMerging(*embed_dim[i: i + 2])) |
| blk.append(nn.Sequential( |
| Residual(Conv2d_BN(embed_dim[i + 1], embed_dim[i + 1], 3, 1, 1, groups=embed_dim[i + 1])), |
| Residual(FFN(embed_dim[i + 1], int(embed_dim[i + 1] * 2))), |
| )) |
|
|
| self.blocks1 = nn.Sequential(*self.blocks1) |
| self.blocks2 = nn.Sequential(*self.blocks2) |
| self.blocks3 = nn.Sequential(*self.blocks3) |
|
|
| self.head = BN_Linear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity() |
| self.distillation = distillation |
| if distillation: |
| self.head_dist = BN_Linear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity() |
|
|
| def forward(self, x): |
| x = self.patch_embed(x) |
| x = self.blocks1(x) |
| x = self.blocks2(x) |
| x = self.blocks3(x) |
| x = nn.functional.adaptive_avg_pool2d(x, 1).flatten(1) |
| if self.distillation: |
| x = self.head(x), self.head_dist(x) |
| if not self.training: |
| x = (x[0] + x[1]) / 2 |
| else: |
| x = self.head(x) |
| return x |
|
|
