Im2Mesh/docs/yaml_configs/train_nnUnet.yaml at main · MMV-Lab/Im2Mesh · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
mode: train-multishape

data:
  category: "pair"
  data_path: "/path/to/training_set" #_IM.tiff (imagen) , _GT.npy (numpy vector) in this case
  extra: # Info for heurisctis nnUnet
    patch_size: [65, 62, 56] # [Z,Y,X]/[Y,X] patch_size ≤ shape of the min volume on the dataset
    spacing: [1.0, 1.0,1.0] # [Z,Y,X]/[Y,X]
    modality: "non-CT"
    min_size: 16
    #vram_gb: 40 #if you use cuda automatically its computed in other case you should provide
  dataloader:
    train:
      dataloader_type:
        module_name: monai.data
        func_name: Dataset # PersistentDataset CacheDataset Dataset
      # dataset_params:
      # # For CacheDataset
      #   # cache_rate: 1.0
      #   # num_workers: 8
      # # For PersistentDataset
      #   # pickle_protocol: 2
      #   # cache_dir: ./tmp0dmp0
      dataloader_params:
        batch_size: 1
        pin_memory: True
        num_workers: 4
        persistent_workers: False
    val:
      dataloader_type:
        module_name: monai.data
        func_name: Dataset # PersistentDataset CacheDataset Dataset
      # dataset_params:
      # # For CacheDataset
      #   # cache_rate: 1.0
      #   # num_workers: 4
      # # For PersistentDataset
      #   # pickle_protocol: 2
      #   # cache_dir: ./tmp0dmp0
      dataloader_params:
        batch_size: 1
        pin_memory: True
        num_workers: 4
        persistent_workers: False
  preprocess:
    - module_name: monai.transforms
      func_name: LoadImaged
      params:
        keys: ["IM"]
        dimension_order_out: "ZYX"
        C: 0
        T: 0
    - module_name: monai.transforms
      func_name: LoadImaged
      params:
        keys: ["GT"]
        dtype: float #int
    - module_name: monai.transforms
      func_name: EnsureChannelFirstd
      params:
        keys: ["IM"]
        channel_dim: "no_channel"
    - module_name: monai.transforms
      func_name: NormalizeIntensityd
      params:
        channel_wise: True
        keys: ["IM"]

    - module_name: mmv_im2im.utils.custom_transforms
      func_name: DivisiblePadWithGTAdjustd
      params:
        keys: ["IM", "GT"]
        k: 16
        mode: "constant"       # zero-padding is safest post-normalisation
        constant_value: 0.0
        n_coord_dims: 3        # number of leading GT elements = spatial coords

    - module_name: monai.transforms
      func_name: EnsureTyped
      params:
        keys: ["IM", "GT"]

  augmentation:
    - module_name: monai.transforms
      func_name: RandHistogramShiftd
      params:
        prob: 0.2
        num_control_points: 50
        keys: ["IM"]

model:
  framework: nnUnet
  model_extra:
    pre-train: '/mnt/eternus/users/Jair/Im2Mesh/lightning_logs/epoch=35-val_loss=0.01130.ckpt'
  #   extend: True  option to tranfer learing (change output layer)
  net:
    module_name: monai.networks.nets
    func_name: DynUNet
    params:
      spatial_dims: 3 # 2d or 3d for convolutions
      in_channels: 1 # nuber of channels in the input IM
      out_channels: 2503 # Vector size on this case [v1,...,vn]-> n
      # DynUNet initial values
      kernel_size: [[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]]
      strides: [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]
      upsample_kernel_size: [[2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]
      filters: [32, 64, 128, 256, 512]
      dropout: 0.2
      res_block: True # Enables residual connections standard in modern nnU-Nets
      deep_supervision: False
      task: 'regression' # segmentation/regression
  criterion:
    module_name: torch.nn
    func_name: HuberLoss #HuberLoss/L1Loss/MSELoss for vectors
    params:
      reduction: 'mean'
      delta: 1.0

  optimizer:
    module_name: torch.optim
    func_name: AdamW
    params:
      lr: 0.01 # 0.001
      weight_decay: 0.01 # 0.01

  scheduler:
    module_name: torch.optim.lr_scheduler
    func_name: ReduceLROnPlateau
    params:
      mode: 'min'
      factor: 0.2
      patience: 25
    monitor: 'val_loss'

trainer:
  verbose: True
  # strategy: ddp_find_unused_parameters_true #avoid timeout
  # gpus: 1 #number or list of gpus to use
  params:
    accumulate_grad_batches: 5
    precision: 32 #16
    max_epochs: 3000
    detect_anomaly: False
    log_every_n_steps: 5 # less than the number of training iterations
    # gradient_clip_val: 0.5
    # gradient_clip_algorithm: "norm"
  callbacks:
    - module_name: lightning.pytorch.callbacks.early_stopping
      func_name: EarlyStopping
      params:
        monitor: 'val_loss'
        patience: 60
        verbose: True
    - module_name: lightning.pytorch.callbacks.model_checkpoint
      func_name: ModelCheckpoint
      params:
        monitor: 'val_loss'
        filename: '{epoch}-{val_loss:.5f}'
        mode: min
        save_top_k: 5
        save_last: true