nnUNet V2代码——nnUNetv2

nnUNet V2代码——nnUNetv2_train命令

完成数据预处理命令后，开始训练

本文目录

训练代码入口
- nnUNetv2_train命令行参数
- run_training函数
训练代码
训练结束

训练代码入口

nnU-Net V2 的训练命令是nnUNetv2_train

nnUNetv2_train命令行参数

参数名称	是否必填	默认值	描述
dataset_name_or_id	是	无	用于训练的数据集名称或 ID
configuration	是	无	需要训练的配置
fold	是	无	5 折交叉验证的折数。应为 0 到 4 之间的整数
-tr	否	nnUNetTrainer	指定自定义训练器。默认值为 nnUNetTrainer
-p	否	nnUNetPlans	指定自定义计划标识符。默认值为 nnUNetPlans
-pretrained_weights	否	None	用于预训练模型的 nnU-Net checkpoint文件路径。仅在实际训练时使用。测试版，请谨慎使用
-num_gpus	否	1	指定训练时使用的 GPU 数量
–use_compressed	否	False	如果设置此标志，训练数据（预处理后生成的压缩版数据集）将不会被解压缩。读取压缩数据会消耗更多 CPU 和（可能）内存，仅在您知道自己在做什么时使用
–npz	否	False	将最终验证的 softmax 预测（数值为概率值，不是类别）保存为 npz 文件（除了预测的分割结果）。这对于找到最佳集成是必需的
–c	否	False	从上次训练结束处开始（已经完成训练则不用）
–val	否	False	设置此标志以仅运行验证。需要训练已完成
–val_best	否	False	如果设置，验证将使用 checkpoint_best 而不是 checkpoint_final。与 --disable_checkpointing 不兼容！警告：这将使用与常规验证相同的“validation”文件夹，无法区分两者！（val数据集没法区分checkpoint_best 还是checkpoint_final）
–disable_checkpointing	否	False	设置此标志以禁用checkpoint保存。适合测试时使用，避免硬盘被checkpoint文件填满
-device	否	cuda	设置训练运行的设备。可用选项为 ‘cuda’（GPU）、‘cpu’（CPU）和 ‘mps’（Apple M1/M2）。不要用此参数设置 GPU ID！请使用 CUDA_VISIBLE_DEVICES=X nnUNetv2_train […] 代替！

确定执行该命令后，首先调用run_training_entry函数，该函数会收集用户在命令行输入的参数，调用同文件下的run_training函数，并将收集的命令行参数传递给它。

run_training_entry函数和run_training函数代码均在nnUNet \ nnunetv2 \ run \ run_training.py文件中。

run_training函数

run_training函数在检查必要参数后，判断GPU数量，多GPU需要配置环境，单GPU不需要。

nnUNetv2_train命令有多处关于多GPU训练的代码，之后会集中一篇文章阅读。🏃🏃🏃

无论哪种情况，run_training函数都有如下操作：

1️⃣首先调用get_trainer_from_args函数，获取用于训练的nnunet_trainer变量，默认是实例化后的nnUNetTrainer类。该函数依次完成查询类、配置文件、实例化，代码结构清晰，不做粘贴：

################# run_training函数部分代码
# 实例化的nnUNetTrainer类
nnunet_trainer = get_trainer_from_args(dataset_name_or_id, configuration, fold, tr, p,
                                        use_compressed)

2️⃣之后完成一些训练前的设置：

################# run_training函数部分代码
# 是否保存网络训练后的权重
if disable_checkpointing:
    nnunet_trainer.disable_checkpointing = disable_checkpointing

assert not (continue_training and only_run_validation), f'Cannot set --c and --val flag at the same time. Dummy.'

# 加载预训练权重
maybe_load_checkpoint(nnunet_trainer, continue_training, only_run_validation, pretrained_weights)

if torch.cuda.is_available():
    cudnn.deterministic = False	# 允许 cuDNN 选择最快的卷积算法，从而加速训练过程
    cudnn.benchmark = True	# 启用 cuDNN 的自动调优功能，找到最适合当前输入大小和硬件的算法，从而加速训练

3️⃣之后运行nnunet_trainer . run_training函数（重名，注意区分 ❗️❗️❗️）和nnunet_trainer . perform_actual_validation函数，完成 train 和 validate：

################# run_training函数部分代码
# 开启训练
if not only_run_validation:
    nnunet_trainer.run_training()

# 是否使用best权重
# nnU-Net V2在训练过程中会生成三个checkpoint.pth
# 分别是checkpoint_best.pth、checkpoint_final.pth、checkpoint_latest.pth
# 由名称可以看出，分别是最佳、最终、最新训练权重
# checkpoint_final.pth会在训练结束时生成，读者如果需要在训练过程中predict，
# 可以在同文件夹下复制checkpoint_best.pth或checkpoint_latest.pth，更改名称后predict
# 上文的参数–val_best也涉及这一点。
if val_with_best:
    nnunet_trainer.load_checkpoint(join(nnunet_trainer.output_folder, 'checkpoint_best.pth'))
# 开始测试val数据集
nnunet_trainer.perform_actual_validation(npz)

训练代码

与预处理命令涉及的类不同，nnUNetTrainer类的各个函数代码较长且嵌套较深，因此，在阅读nnUNetTrainer类时，我们不再采用单个函数一个接一个的方式阅读代码，而是按照训练过程依次阅读，以便更好地理解其整体流程和设计思路。

根据上文，训练过程主要由nnunet_trainer变量的run_training函数（重名，注意区分 ❗️❗️❗️）完成，该变量默认是实例化的nnUNetTrainer类（nnU-Net V2文档介绍到，读者可以自定义该类）。

以下将用run_training函数指称nnunet_trainer . run_training函数，不再和上文的run_training函数区分

run_training函数代码如下：

######################## run_training函数代码
def run_training(self):
	### 训练开始
    self.on_train_start()

    for epoch in range(self.current_epoch, self.num_epochs):
    	### epoch开始
        self.on_epoch_start()
        
		### epoch train 开始
        self.on_train_epoch_start()
        train_outputs = []
       	### 一个epoch会train 250次（默认值，在nnUNetTrainer类的__init__函数中会讲到）
        for batch_id in range(self.num_iterations_per_epoch):
        	### 250 次的一次，one step
            train_outputs.append(self.train_step(next(self.dataloader_train)))
        ### epoch train 结束
        self.on_train_epoch_end(train_outputs)

        with torch.no_grad():
        	### epoch val 开始
            self.on_validation_epoch_start()
            val_outputs = []
            ### 一个epoch会val 50次（默认值，在nnUNetTrainer类的__init__函数中会讲到）
            for batch_id in range(self.num_val_iterations_per_epoch):
            	### 50 次的一次，one step
                val_outputs.append(self.validation_step(next(self.dataloader_val)))
        	### epoch val 结束
            self.on_validation_epoch_end(val_outputs)
            
		### epoch结束
        self.on_epoch_end()
	### 训练结束
    self.on_train_end()