如何修改预训练PyTorch模型最后一层以获取倒数第二层输出？

方法1：替换最后一层线性层为恒等映射 #

这是最直接的方案，和ResNet18的思路一致——找到模型中最后那个nn.Linear分类层，将其替换为nn.Identity()（该层会直接返回输入，不做任何变换）。

操作步骤：

1. 加载训练好的模型：

import torch
import torch.nn as nn

# 假设你的模型类定义如下
class MyModel(nn.Module):
    def __init__(self, num_classes=45):
        super().__init__()
        # 特征提取模块，最终输出512维特征
        self.features = nn.Sequential(
            # 这里是你的卷积、池化等特征提取层
        )
        self.fc = nn.Linear(512, num_classes) # 最后一层分类线性层

    def forward(self, x):
        x = self.features(x)
        x = self.fc(x)
        return x

# 加载训练好的权重文件
model = MyModel()
model.load_state_dict(torch.load('your_trained_model_weights.pth'))

2. 替换最后一层：
   如果你的最后一层属性名是fc（和示例一致），直接执行：

model.fc = nn.Identity()

如果最后一层叫其他名字（比如classifier），先通过以下代码确认层名：

for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        print(f"线性层名称：{name}，结构：{module}")

找到对应名称后替换即可，比如：

model.classifier = nn.Identity()

3. 获取特征向量：
   此时输入数据到模型，输出就是512维的特征向量：

# 示例输入（根据你的模型输入尺寸调整）
input_tensor = torch.randn(1, 3, 224, 224)
feature_vector = model(input_tensor)
print(feature_vector.shape) # 输出: torch.Size([1, 512])

方法2：修改前向传播逻辑返回特征 #

如果你不想修改模型结构，可以直接修改模型的forward方法，让它返回倒数第二层的特征，而非分类结果。

修改后的模型类：

class MyModel(nn.Module):
    def __init__(self, num_classes=45):
        super().__init__()
        self.features = nn.Sequential(
            # 特征提取层
        )
        self.fc = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.features(x) # 这一步输出的就是512维特征
        # 注释掉分类层的计算，直接返回特征
        # x = self.fc(x)
        return x

加载权重后直接调用模型即可得到特征向量。

方法3：临时提取特征（不修改模型） #

如果不想改动模型定义或结构，可以写一个辅助函数，手动遍历模型层获取倒数第二层输出：

def extract_features(model, x):
    # 遍历模型子模块，跳过最后一层分类层
    for name, module in model.named_children():
        if name != 'fc': # 替换成你的最后一层名称
            x = module(x)
    return x

# 使用示例
feature_vector = extract_features(model, input_tensor)

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