如何从MAE模型输出提取预测标签并计算10类混淆矩阵？

1. 正确提取预测标签 #

你的模型输出里，outputs[1]（即PRED）是预测相关张量，形状为[1, 196, 3328]。结合10类图像分类场景，这里的3328是每个图像patch的特征维度，需要先对这些patch特征做全局聚合，再通过分类头得到对应10类的logits，最后取argmax得到预测标签。

示例代码片段：

import torch

# 拆分模型输出
loss, pred, mask, target = outputs

# 对patch特征做全局平均池化，得到单样本全局特征
global_features = torch.mean(pred, dim=1)  # 形状变为[1, 3328]

# 传入分类头生成10类logits（假设分类头为model.classifier）
logits = model.classifier(global_features)  # 形状变为[1, 10]

# 提取预测标签（取logits最大值对应的类别索引）
pred_labels = torch.argmax(logits, dim=1).cpu().numpy()

# 提取真实标签（注意：分类任务的真实标签来自数据加载器的输入，而非模型输出的target）
true_labels = data_labels.cpu().numpy()  # data_labels是当前batch的输入标签

如果你的模型已在输出中直接给出分类logits（若形状应为[batch_size, 10]，需核对PRED的实际含义），可简化为：

pred_labels = torch.argmax(pred, dim=1).cpu().numpy()
true_labels = data_labels.cpu().numpy()

2. 计算混淆矩阵 #

收集所有批次的真实标签与预测标签后，用sklearn工具计算混淆矩阵：

from sklearn.metrics import confusion_matrix
import numpy as np

# 初始化容器存储所有样本标签
all_true = []
all_pred = []

# 遍历数据集循环
for inputs, data_labels in dataloader:
    outputs = model(inputs)
    loss, pred, mask, target = outputs
    
    # 提取当前批次标签
    global_features = torch.mean(pred, dim=1)
    logits = model.classifier(global_features)
    pred_labels = torch.argmax(logits, dim=1).cpu().numpy()
    true_labels = data_labels.cpu().numpy()
    
    all_true.extend(true_labels)
    all_pred.extend(pred_labels)

# 生成混淆矩阵
cm = confusion_matrix(all_true, all_pred)
print("混淆矩阵：")
print(cm)

关键注意事项 #

• 区分target的用途：MAE预训练阶段的target是图像patch重构目标，分类任务的真实标签必须来自数据加载器的输入data_labels，不能用模型输出的outputs[3]
• 若分类模型已直接输出logits，需确认PRED的实际维度，避免对patch特征做多余处理

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