Im2Mesh/docs/yaml_configs/train_Attention_Unet.yaml at main · JairMathAI/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
mode: train-multishape

data:
  category: "pair"
  data_path: "/path/to/training_set" #_IM.tiff (imagen) , _GT.npy (numpy vector) in this case
  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: FCN
  # model_extra:
  #   pre-train: path to ckp file
  #   extend: True  option to tranfer learing (change output layer)
  net:
    module_name: monai.networks.nets
    func_name: AttentionUnet
    params:
      in_channels: 1 # nuber of channels in the input IM
      out_channels: 2503 # Vector size on this case [v1,...,vn]-> n
      spatial_dims: 3 # 2d or 3d for convolutions
      channels: [32, 64, 128, 256, 512]
      strides: [1, 2, 2, 2, 2]
      dropout: 0.2
      task: "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 1 in case of accumulate grad
    # 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