如何从torchvision的VisionTransformer(ViT)提取768维特征并解决报错

错误原因 #

你仿照CNN移除输出层的方式处理ViT，但忽略了ViT的内部运行逻辑：ViT的conv_proj（patch嵌入层）输出是(批量大小, 768, 14, 14)的张量，而后续的encoder需要接收序列格式的输入（(批量大小, 序列长度, 768)），其中序列长度包含分类token和所有patch token（共197个）。直接用nn.Sequential拼接conv_proj和encoder，会跳过维度转换、分类token添加、位置编码这些关键步骤，导致维度不匹配报错。

正确实现方式 #

方法一：修改模型分类头（最简单）

直接将ViT的分类头替换为恒等映射，模型会直接返回分类token对应的768维特征：

from torch import nn
from torchvision.models.vision_transformer import vit_b_16
from torchvision.models import ViT_B_16_Weights
from PIL import Image as PIL_Image

# 加载预训练ViT模型
vit = vit_b_16(weights=ViT_B_16_Weights.DEFAULT)
# 替换分类头为恒等映射，使模型输出cls token的特征
vit.heads = nn.Identity()

# 获取官方预处理流程
preprocessing = ViT_B_16_Weights.DEFAULT.transforms()

# 加载并预处理图像
img = PIL_Image.open("example.png")
img_tensor = preprocessing(img).unsqueeze(0)  # 增加批量维度

# 提取768维特征
features = vit(img_tensor)
print(features.shape)  # 输出: torch.Size([1, 768])

方法二：手动处理特征流程（更灵活）

如果需要自定义特征提取逻辑（比如提取所有patch的特征），可以手动复现ViT的forward步骤：

import torch
from torch import nn
from torchvision.models.vision_transformer import vit_b_16
from torchvision.models import ViT_B_16_Weights
from PIL import Image as PIL_Image

vit = vit_b_16(weights=ViT_B_16_Weights.DEFAULT)
# 提取ViT核心组件
conv_proj = vit.conv_proj
encoder = vit.encoder
class_token = vit.class_token
pos_embedding = vit.pos_embedding

preprocessing = ViT_B_16_Weights.DEFAULT.transforms()
img = PIL_Image.open("example.png")
img_tensor = preprocessing(img).unsqueeze(0)  # shape: (1, 3, 224, 224)

# 1. 生成patch嵌入
x = conv_proj(img_tensor)  # shape: (1, 768, 14, 14)
x = x.flatten(2).transpose(1, 2)  # 转换为序列格式: (1, 196, 768)

# 2. 添加分类token
batch_class_token = class_token.expand(img_tensor.shape[0], -1, -1)  # shape: (1, 1, 768)
x = torch.cat([batch_class_token, x], dim=1)  # shape: (1, 197, 768)

# 3. 添加位置编码
x = x + pos_embedding

# 4. 经过encoder提取特征
x = encoder(x)

# 5. 提取分类token对应的768维特征
features = x[:, 0]  # shape: (1, 768)

内容的提问来源于stack exchange，提问作者mitja