代码复现记录nnUNet

1
# 基本训练命令
2
nnUNetv2_train DATASET_ID CONFIG FOLD [其他参数]
3

4
# 一些可用的命令行参数：
5
nnUNetv2_train 001 2d 0 --npz  # 使用npz格式保存
6
nnUNetv2_train 001 2d 0 --c    # 继续之前的训练
7
nnUNetv2_train 001 2d 0 --val  # 同时进行验证

1
def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
2
             device: torch.device = torch.device('cuda')):
3
    # 修改学习率
4
    self.initial_lr = 1e-2  # 默认是1e-2，可以修改为其他值
5

6
    # 修改批次大小会在后面的方法中设置
7

8
    # 修改训练轮数
9
    self.num_epochs = 1000  # 默认1000个epoch

1
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
2

3
class MyCustomTrainer(nnUNetTrainer):
4
    def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, unpack_dataset: bool = True,
5
                 device: torch.device = torch.device('cuda')):
6
        super().__init__(plans, configuration, fold, dataset_json, unpack_dataset, device)
7

8
        # 自定义学习率
9
        self.initial_lr = 5e-3  # 修改初始学习率
10

11
        # 自定义训练轮数
12
        self.num_epochs = 500
13

14
        # 自定义权重衰减
15
        self.weight_decay = 3e-5
16

17
    def configure_optimizers(self):
18
        # 自定义优化器参数
19
        optimizer = torch.optim.SGD(
20
            self.network.parameters(),
21
            lr=self.initial_lr,
22
            weight_decay=self.weight_decay,
23
            momentum=0.99,
24
            nesterov=True
25
        )
26
        return optimizer
27

28
    def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
29
        # 修改数据增强参数
30
        rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \
31
            super().configure_rotation_dummyDA_mirroring_and_inital_patch_size()
32

33
        # 可以在这里修改patch size等参数
34
        return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes

1
# 在自定义trainer中重写方法
2
def configure_optimizers(self):
3
    # 强制设置batch size
4
    self.batch_size = 4  # 设置你想要的batch size
5

6
    # 重新计算相关参数
7
    optimizer = torch.optim.SGD(
8
        self.network.parameters(),
9
        lr=self.initial_lr,
10
        weight_decay=self.weight_decay,
11
        momentum=0.99,
12
        nesterov=True
13
    )
14
    return optimizer

1
# 使用自定义trainer训练
2
nnUNetv2_train DATASET_ID CONFIG FOLD -tr MyCustomTrainer

1
class MyCustomTrainer(nnUNetTrainer):
2
    def __init__(self, *args, **kwargs):
3
        super().__init__(*args, **kwargs)
4

5
        # 学习率相关
6
        self.initial_lr = 1e-3           # 初始学习率
7
        self.lr_scheduler_eps = 1e-8     # 学习率调度器参数
8
        self.lr_scheduler_patience = 30  # 学习率调度器耐心值
9

10
        # 训练相关
11
        self.num_epochs = 500            # 训练轮数
12
        self.num_iterations_per_epoch = 250  # 每个epoch的迭代次数
13
        self.save_every = 25             # 每多少个epoch保存一次
14
        self.validate_every = 5          # 每多少个epoch验证一次
15

16
        # 优化器相关
17
        self.weight_decay = 3e-5         # 权重衰减
18

19
        # 损失函数相关参数
20
        self.deep_supervision_scales = [1, 0.5, 0.25, 0.125]  # 深度监督的尺度

1
export nnUNet_n_proc_DA=8  # 数据增强进程数
2
nnUNetv2_train 001 2d 0