如何提取PyTorch中Vision Transformer的倒数第二层输出？

我用DataParallel在自定义数据集上微调了Vision Transformer模型，代码如下：

import torch
import torch.nn as nn
import timm

model = timm.create_model('vit_base_patch16_224', pretrained=False)
model.head = nn.Sequential(
    nn.Linear(768, 512),
    nn.ReLU(),
    nn.BatchNorm1d(512),
    nn.Dropout(p=0.2),
    nn.Linear(512, 141)
)
checkpoint = torch.load('vit_b_16v3.pth')
# 移除checkpoint中的module.前缀（适配DataParallel保存的权重）
checkpoint = {k.partition('module.')[2]: v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint)

模型最后几层结构：

(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
(fc_norm): Identity()
(head_drop): Dropout(p=0.0, inplace=False)
(head): Sequential(
  (0): Linear(in_features=768, out_features=512, bias=True)
  (1): ReLU()
  (2): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (3): Dropout(p=0.2, inplace=False)
  (4): Linear(in_features=512, out_features=141, bias=True)
)

需要输入图像后得到512维的描述向量（对应最后一个Linear层之前的输出），尝试过一些提取中间层的方法但无效，以下是可行方案：

方法1：修改模型forward方法 #

重写模型的forward函数，同时返回分类结果和512维特征：

class FeatureExtractingViT(timm.models.vision_transformer.VisionTransformer):
    def forward(self, x):
        # 保留原ViT前向逻辑
        x = self.patch_embed(x)
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
        if self.dist_token is None:
            x = torch.cat((cls_token, x), dim=1)
        else:
            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
        x = self.pos_drop(x + self.pos_embed)
        x = self.blocks(x)
        x = self.norm(x)
        x = x[:, 0] if self.dist_token is None else x[:, 0] + x[:, 1]
        
        # 提取head前4层的输出（512维特征）
        feat = x
        for layer in self.head[:4]:
            feat = layer(feat)
        # 计算最终分类结果
        logits = self.head[4](feat)
        
        return logits, feat

替换原模型并推理：

# 初始化模型并加载权重
model = FeatureExtractingViT('vit_base_patch16_224', pretrained=False)
model.head = nn.Sequential(
    nn.Linear(768, 512),
    nn.ReLU(),
    nn.BatchNorm1d(512),
    nn.Dropout(p=0.2),
    nn.Linear(512, 141)
)
checkpoint = torch.load('vit_b_16v3.pth')
checkpoint = {k.partition('module.')[2]: v for k, v in checkpoint.items()}
model.load_state_dict(checkpoint)

# 推理获取特征
model.eval()
with torch.no_grad():
    input_img = torch.randn(1, 3, 224, 224)  # 示例输入
    logits, feat_512 = model(input_img)
    print(feat_512.shape)  # 输出: torch.Size([1, 512])

方法2：使用正向钩子捕获中间层 #

无需修改模型结构，注册钩子获取head[3]（Dropout层）的输出：

# 定义钩子函数存储特征
feature = None
def hook_fn(module, input, output):
    global feature
    feature = output

# 为head的第3层注册钩子
hook = model.head[3].register_forward_hook(hook_fn)

# 推理
model.eval()
with torch.no_grad():
    input_img = torch.randn(1, 3, 224, 224)
    logits = model(input_img)
    print(feature.shape)  # 输出: torch.Size([1, 512])

# 用完钩子后移除
hook.remove()

方法3：构建特征提取子模型 #

截取模型到倒数第二层，创建独立的特征提取器：

class ClsTokenExtractor(nn.Module):
    def __init__(self, dist_token=None):
        super().__init__()
        self.dist_token = dist_token
    
    def forward(self, x):
        return x[:, 0] if self.dist_token is None else x[:, 0] + x[:, 1]

# 构建特征提取器
feature_extractor = nn.Sequential(
    model.patch_embed,
    model.pos_drop,
    model.blocks,
    model.norm,
    ClsTokenExtractor(model.dist_token),
    model.head[:4]
)

# 推理
feature_extractor.eval()
with torch.no_grad():
    input_img = torch.randn(1, 3, 224, 224)
    feat_512 = feature_extractor(input_img)
    print(feat_512.shape)  # 输出: torch.Size([1, 512])

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