### 高效视觉Transformer (ViT) 的代码实现 以下是基于 `EfficientViT` 和传统 ViT 设计理念的 Python 实现示例。此代码展示了如何通过分组注意力机制和优化策略提升模型效率。 #### 多头自注意力 (MHSA) 层 多头自注意力层是 ViT 中的核心组件，负责捕捉全局特征关系[^1]： ```python import torch import torch.nn as nn class MultiHeadSelfAttention(nn.Module): def __init__(self, dim, num_heads=8, dropout=0.0): super(MultiHeadSelfAttention, self).__init__() assert dim % num_heads == 0, "dim should be divisible by num_heads." self.num_heads = num_heads self.head_dim = dim // num_heads self.qkv = nn.Linear(dim, dim * 3) self.dropout = nn.Dropout(dropout) self.proj = nn.Linear(dim, dim) def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv.unbind(0) attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5) attn = attn.softmax(dim=-1) attn = self.dropout(attn) out = (attn @ v).transpose(1, 2).reshape(B, N, C) out = self.proj(out) return out ``` #### EfficientViT 主体结构 为了进一步减少内存消耗并提高计算效率，可以引入级联分组注意力机制[^3]： ```python class CascadedGroupAttention(nn.Module): def __init__(self, dim, group_size=8, num_heads=4, dropout=0.0): super(CascadedGroupAttention, self).__init__() self.group_attention = MultiHeadSelfAttention(dim, num_heads=num_heads, dropout=dropout) self.group_size = group_size def forward(self, x): B, N, C = x.shape grouped_x = x.view(B, -1, self.group_size, C) # 分组操作 attended_grouped_x = self.group_attention(grouped_x.reshape(B * (N // self.group_size), self.group_size, C)) output = attended_grouped_x.view(B, N, C) return output class EfficientViTBlock(nn.Module): def __init__(self, dim, mlp_ratio=4., drop_rate=0.0, group_size=8, num_heads=4): super(EfficientViTBlock, self).__init__() self.norm1 = nn.LayerNorm(dim) self.attn = CascadedGroupAttention(dim, group_size=group_size, num_heads=num_heads, dropout=drop_rate) self.drop_path = DropPath(drop_rate) if drop_rate > 0. else nn.Identity() self.norm2 = nn.LayerNorm(dim) hidden_features = int(mlp_ratio * dim) self.mlp = Mlp(in_features=dim, hidden_features=hidden_features, act_layer=nn.GELU, drop=drop_rate) def forward(self, x): x = x + self.drop_path(self.attn(self.norm1(x))) x = x + self.drop_path(self.mlp(self.norm2(x))) return x ``` #### 完整模型定义 最后，构建完整的 EfficientViT 模型框架[^2]: ```python class EfficientViT(nn.Module): def __init__(self, image_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., group_size=8, drop_rate=0.0, num_classes=1000): super(EfficientViT, self).__init__() self.patch_embed = PatchEmbed(image_size=image_size, patch_size=patch_size, in_chans=3, embed_dim=embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) self.blocks = nn.Sequential(*[ EfficientViTBlock(dim=embed_dim, mlp_ratio=mlp_ratio, drop_rate=drop_rate, group_size=group_size, num_heads=num_heads) for _ in range(depth)]) self.norm = nn.LayerNorm(embed_dim) self.fc = nn.Linear(embed_dim, num_classes) def forward(self, x): x = self.patch_embed(x) cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_tokens, x), dim=1) x += self.pos_embed x = self.blocks(x) x = self.norm(x[:, 0]) x = self.fc(x) return x ``` --- ###