如何提取PyTorch训练的Autoencoder隐藏层特征用于后续分类任务

你定义的AutoEncoder结构中，encoder层的输出就是需要的2维隐藏特征，直接调用encoder模块或取forward返回的第一个值即可完成提取，具体实现如下：

1. 预操作：加载训练好的模型 #

首先加载已保存的模型，切换到评估模式，避免推理阶段修改模型参数：

import torch
import numpy as np

# 加载训练完成的模型
model = torch.load('autoencoder.pth')
# 切换为评估模式
model.eval()

2. 隐藏特征提取实现 #

提供两种常用的提取方式，可按需选择：

方式1：提取全量数据集的隐藏特征 #

适合需要一次性获取所有样本特征的场景：

# 输入数据格式和你预处理阶段的逻辑保持一致
all_input = torch.Tensor(input.to_numpy())

feature_list = []
# 关闭梯度计算，节省显存、提升运行速度
with torch.no_grad():
    # 直接调用encoder模块得到隐藏特征
    features = model.encoder(all_input.view(-1, 84))
    # 转为numpy数组，可直接输入后续分类算法使用
    feature_array = features.cpu().numpy()

# 输出特征维度，正常应为 (总样本数, 2)
print(feature_array.shape)

方式2：分别提取训练集、测试集的隐藏特征 #

适合需要和之前的数据集拆分逻辑对齐、按批次处理的场景：

def extract_features(dataloader):
    all_features = []
    all_labels = []
    with torch.no_grad():
        for data, labels in dataloader:
            inputs = data.view(-1, 84)
            # 取forward返回的第一个值即为隐藏特征
            codes, _ = model(inputs)
            all_features.append(codes.cpu().numpy())
            all_labels.append(labels.cpu().numpy())
    # 拼接所有批次的结果
    return np.concatenate(all_features), np.concatenate(all_labels)

# 提取训练集特征和对应标签
train_features, train_labels = extract_features(train_loader)
# 提取测试集特征和对应标签
test_features, test_labels = extract_features(test_loader)

注意事项 #

• 提取特征前必须调用model.eval()，否则模型会保持训练模式，输出结果不符合预期
• 务必使用with torch.no_grad()包裹提取逻辑，避免梯度计算占用不必要的内存资源

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