在nnUNet中集成WandB进行训练可视化与日志记录

概述

nnUNet作为医学图像分割领域的标杆框架，其训练过程的可视化对于模型调优和结果分析至关重要。本文将详细介绍如何在nnUNet框架中集成Weights & Biases(WandB)工具，实现训练过程的实时监控、指标记录和图像可视化。

WandB集成方案

准备工作

首先需要安装WandB客户端库：

pip install wandb

自定义Trainer实现

nnUNet采用模块化设计，最佳实践是通过继承基础Trainer类来实现自定义功能：

1. 在nnunetv2/training/nnUNetTrainer目录下创建新的训练器类
2. 重写关键训练方法，插入WandB日志记录逻辑

import wandb
import numpy as np
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer

class nnUNetTrainerWandB(nnUNetTrainer):
    def __init__(self, plans, configuration, fold, output_folder):
        super().__init__(plans, configuration, fold, output_folder)
        wandb.init(project="nnUNet_Medical_Segmentation")
        
    def on_epoch_end(self):
        super().on_epoch_end()
        # 记录训练指标
        wandb.log({
            "train_loss": self.logger.my_fantastic_loss,
            "val_loss": self.logger.my_fantastic_val_loss,
            "epoch": self.current_epoch
        })
        
    def validate(self, *args, **kwargs):
        outputs = super().validate(*args, **kwargs)
        # 可视化验证结果
        for batch in validation_loader:
            wandb.log({
                "input_image": wandb.Image(batch['image'][0]),
                "prediction": wandb.Image(outputs[0].argmax(0)),
                "ground_truth": wandb.Image(batch['label'][0])
            })
        return outputs

关键实现细节

1. 图像数据预处理：医学图像通常需要特殊处理才能正确显示

   • 窗宽窗位调整
   • 多通道转RGB
   • 归一化处理

2. 指标选择：除基础loss外，建议记录：

   • Dice系数
   • Hausdorff距离
   • 特定器官的分割指标

3. 采样策略：全量验证集可视化会导致日志过大，建议：

   • 每N个epoch记录一次
   • 随机选择部分切片展示
   • 使用低分辨率预览

训练执行

使用自定义Trainer启动训练：

nnUNetv2_train [...] -tr nnUNetTrainerWandB

高级技巧

1. 超参数记录：自动记录所有训练配置参数

   wandb.config.update(self.plans)
   

2. 模型检查点保存：将最佳模型保存到WandB

   wandb.save(os.path.join(self.output_folder, "*.pth"))
   

3. 资源监控：记录GPU显存、CPU利用率等系统指标

可视化效果优化

1. 多模态融合显示：对于多模态数据，可采用叠加显示
2. 轮廓对比：预测结果与GT以轮廓形式叠加显示
3. 错误区域高亮：用不同颜色标记FP/FN区域
4. 3D切片浏览：上传多张切片生成可交互浏览面板

常见问题处理

1. 数据保密问题：

   • 使用WandB的私有项目
   • 上传前匿名化处理
   • 仅显示ROI区域

2. 大图像处理：

   • 下采样显示
   • 切片展示
   • 使用WandB的媒体面板

3. 指标异常检测：

   • 设置合理阈值告警
   • 记录梯度分布
   • 跟踪参数变化

通过以上方法，研究人员可以全面掌握nnUNet模型的训练动态，及时发现并解决问题，大幅提高医学图像分割模型的开发效率。