如何将PyTorch Lightning中torchmetrics生成的混淆矩阵记录到Wandb？

针对你的需求，有两种可行方案：

方案一：直接将已生成的混淆矩阵转为Wandb可视化热图 #

不需要还原原始标签和预测结果，Wandb支持直接基于已计算的混淆矩阵生成可视化。修改你的validation_epoch_end方法即可：

def validation_epoch_end(self, outputs):
    confusion_matrix = self.matrix.compute().cpu().numpy()
    class_names = [f"Class {i}" for i in range(9)]  # 替换为你的实际类别名称
    
    # 方法1：用Wandb内置混淆矩阵绘图工具
    conf_mat_plot = wandb.plot.confusion_matrix(
        probs=None,
        y_true=None,
        preds=None,
        class_names=class_names,
        matrix=confusion_matrix
    )
    wandb.log({"confusion_matrix": conf_mat_plot})
    
    # 方法2：用Seaborn绘制热图后转为Wandb图片
    import seaborn as sns
    import matplotlib.pyplot as plt
    plt.figure(figsize=(10, 8))
    sns.heatmap(confusion_matrix, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('预测类别')
    plt.ylabel('真实类别')
    plt.title('验证集混淆矩阵')
    wandb.log({"confusion_matrix_heatmap": wandb.Image(plt)})
    plt.close()
    
    self.matrix.reset()  # 重置混淆矩阵，避免跨epoch数据污染

方案二：收集真实标签与预测结果，传给Wandb原生记录器 #

如果你偏好使用Wandb原生的混淆矩阵记录逻辑，可以在验证阶段收集所有批次的标签和预测结果，再统一传入：

def _shared_eval_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    acc = accuracy(y_hat, y)
    self.matrix.update(y_hat, y)
    
    # 收集真实标签和预测类别
    y_pred = torch.argmax(y_hat, dim=1)
    self.y_trues.append(y.cpu())
    self.y_preds.append(y_pred.cpu())
    return loss, acc

def validation_epoch_end(self, outputs):
    # 合并所有批次的标签和预测
    y_trues = torch.cat(self.y_trues).numpy()
    y_preds = torch.cat(self.y_preds).numpy()
    
    # 用Wandb原生方法记录混淆矩阵
    wandb.log({
        "confusion_matrix": wandb.plot.confusion_matrix(
            y_true=y_trues,
            preds=y_preds,
            class_names=[f"Class {i}" for i in range(9)]  # 替换为你的实际类别名称
        )
    })
    
    # 重置收集容器和混淆矩阵
    self.y_trues.clear()
    self.y_preds.clear()
    self.matrix.reset()

注意事项 #

• 每次epoch结束后必须重置self.matrix、self.y_trues和self.y_preds，防止跨epoch的数据累积导致错误。
• 如果你的数据集有自定义类别名称，替换class_names列表为实际名称，能让可视化更直观。

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