torch ViT提取特征

### ViT 特征提取方法 为了使用 PyTorch 实现 Vision Transformer (ViT) 进行图像特征提取，可以基于已有的实现代码构建模型，并调整其输出层以获取中间表示。以下是具体的方法： #### 构建 ViT 模型 首先定义 ViT 模型结构，通常包括以下几个部分： - 图像分块（Patch Embedding） - 位置编码（Positional Encoding） - 多头自注意力机制（Multi-head Self-Attention） - 前馈网络（Feed Forward Network） 以下是一个简化版的 ViT 模型实现[^2]： ```python import torch from torch import nn class PatchEmbed(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): super().__init__() self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): x = self.proj(x).flatten(2).transpose(1, 2) return x class ViT(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim) self.pos_embed = nn.Parameter(torch.zeros(1, (img_size // patch_size)**2 + 1, embed_dim)) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.blocks = nn.Sequential(*[ Block(dim=embed_dim, num_heads=num_heads) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) def forward(self, x): batch_size = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(batch_size, -1, -1) x = torch.cat((cls_tokens, x), dim=1) x += self.pos_embed[:, :x.size(1)] x = self.blocks(x) x = self.norm(x) return x ``` 上述代码实现了基本的 ViT 结构，其中 `Block` 是多头自注意力和前馈网络的具体实现。 --- #### 提取特征 要从 ViT 中提取图像特征，可以通过截断模型的方式保留特定层数的输出。例如，在最后一层 Transformer 输出之后但分类头之前获取特征向量。 修改后的前向传播函数如下所示： ```python def extract_features(self, x): batch_size = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(batch_size, -1, -1) x = torch.cat((cls_tokens, x), dim=1) x += self.pos_embed[:, :x.size(1)] x = self.blocks(x) features = self.norm(x) # 获取标准化后的所有补丁特征 return features ``` 调用此函数即可得到输入图像对应的特征矩阵。 --- #### 示例代码 完整的特征提取流程可参考以下代码片段： ```python # 初始化模型 model = ViT() # 输入张量模拟一张图片 input_tensor = torch.randn(1, 3, 224, 224) # 提取特征 features = model.extract_features(input_tensor) print(f"Features shape: {features.shape}") # 输出形状应为 [batch_size, tokens_num, embed_dim] ``` 此处 `tokens_num` 表示图像被划分为多少个补丁加上一个 CLS token，而 `embed_dim` 则是嵌入维度大小。 --- #### 注意事项 1. **预训练权重加载** 如果希望提升效果，建议加载官方或第三方提供的预训练权重文件[^1]。 2. **数据预处理** 确保输入图像经过适当缩放至指定尺寸（如 \(224 \times 224\)），并完成归一化操作。 3. **硬件需求** ViT 对显存消耗较大，推荐在 GPU 上运行以加速计算过程。 ---