complete the model package

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+models/model.ts filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,79 @@
+---
+tags:
+- monai
+- medical
+library_name: monai
+license: unknown
+---
+# Description
+A neural architecture search algorithm for volumetric (3D) segmentation of the pancreas and pancreatic tumor from CT image.
+
+# Model Overview
+This model is trained using the state-of-the-art algorithm [1] of the "Medical Segmentation Decathlon Challenge 2018" with 196 training images, 56 validation images, and 28 testing images.
+
+## Data
+The training dataset is Task07_Pancreas.tar from http://medicaldecathlon.com/. And the data list/split can be created with the script `scripts/prepare_datalist.py`.
+
+## Training configuration
+The training was performed with at least 16GB-memory GPUs.
+
+Actual Model Input: 96 x 96 x 96
+
+## Input and output formats
+Input: 1 channel CT image
+
+Output: 3 channels: Label 2: pancreatic tumor; Label 1: pancreas; Label 0: everything else
+
+## Scores
+This model achieves the following Dice score on the validation data (our own split from the training dataset):
+
+Mean Dice = 0.72
+
+## commands example
+Create data split (.json file):
+
+```
+python scripts/prepare_datalist.py --path /path-to-Task07_Pancreas/ --output configs/dataset_0.json
+```
+
+Execute model searching:
+
+```
+python -m scripts.search run --config_file configs/search.yaml
+```
+
+Execute multi-GPU model searching (recommended):
+
+```
+torchrun --nnodes=1 --nproc_per_node=8 -m scripts.search run --config_file configs/search.yaml
+```
+
+Execute training:
+
+```
+python -m monai.bundle run training --meta_file configs/metadata.json --config_file configs/train.yaml --logging_file configs/logging.conf
+```
+
+Override the `train` config to execute multi-GPU training:
+
+```
+torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run training --meta_file configs/metadata.json --config_file "['configs/train.yaml','configs/multi_gpu_train.yaml']" --logging_file configs/logging.conf
+```
+
+Override the `train` config to execute evaluation with the trained model:
+
+```
+python -m monai.bundle run evaluating --meta_file configs/metadata.json --config_file "['configs/train.yaml','configs/evaluate.yaml']" --logging_file configs/logging.conf
+```
+
+Execute inference:
+
+```
+python -m monai.bundle run evaluating --meta_file configs/metadata.json --config_file configs/inference.yaml --logging_file configs/logging.conf
+```
+
+# Disclaimer
+This is an example, not to be used for diagnostic purposes.
+
+# References
+[1] He, Y., Yang, D., Roth, H., Zhao, C. and Xu, D., 2021. Dints: Differentiable neural network topology search for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 5841-5850).

configs/evaluate.yaml

@@ -0,0 +1,64 @@
+---
+validate#postprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: Activationsd
+    keys: pred
+    softmax: true
+  - _target_: Invertd
+    keys:
+    - pred
+    - label
+    transform: "@validate#preprocessing"
+    orig_keys: image
+    meta_key_postfix: meta_dict
+    nearest_interp:
+    - false
+    - true
+    to_tensor: true
+  - _target_: AsDiscreted
+    keys:
+    - pred
+    - label
+    argmax:
+    - true
+    - false
+    to_onehot: 3
+  - _target_: CopyItemsd
+    keys: "pred"
+    times: 1
+    names: "pred_save"
+  - _target_: AsDiscreted
+    keys:
+    - pred_save
+    argmax:
+    - true
+  - _target_: SaveImaged
+    keys: pred_save
+    meta_keys: pred_meta_dict
+    output_dir: "@output_dir"
+    resample: false
+    squeeze_end_dims: true
+validate#dataset:
+    _target_: Dataset
+    data: "@val_datalist"
+    transform: "@validate#preprocessing"
+validate#handlers:
+- _target_: CheckpointLoader
+  load_path: "$@ckpt_dir + '/model.pt'"
+  load_dict:
+    model: "@network"
+- _target_: StatsHandler
+  iteration_log: false
+- _target_: MetricsSaver
+  save_dir: "@output_dir"
+  metrics:
+  - val_mean_dice
+  - val_acc
+  metric_details:
+  - val_mean_dice
+  batch_transform: "$monai.handlers.from_engine(['image_meta_dict'])"
+  summary_ops: "*"
+evaluating:
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+- "$@validate#evaluator.run()"

configs/inference.yaml

@@ -0,0 +1,116 @@
+---
+imports:
+- "$import glob"
+- "$import os"
+input_channels: 1
+output_classes: 3
+arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+bundle_root: "/workspace/MONAI/model-zoo/models/pancreas_ct_dints_segmentation"
+output_dir: "$@bundle_root + '/eval'"
+dataset_dir: "/workspace/data/msd/Task07_Pancreas"
+data_list_file_path: "$@bundle_root + '/configs/dataset_0.json'"
+datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='testing',
+  base_dir=@dataset_dir)"
+device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
+dints_space:
+  _target_: monai.networks.nets.TopologyInstance
+  channel_mul: 1
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  arch_code:
+  - "$@arch_ckpt['arch_code_a']"
+  - "$@arch_ckpt['arch_code_c']"
+  device: "$torch.device('cuda')"
+network_def:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true
+  node_a: "$torch.from_numpy(@arch_ckpt['node_a'])"
+network: "$@network_def.to(@device)"
+preprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys: image
+  - _target_: EnsureChannelFirstd
+    keys: image
+  - _target_: Orientationd
+    keys: image
+    axcodes: RAS
+  - _target_: Spacingd
+    keys: image
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode: bilinear
+  - _target_: ScaleIntensityRanged
+    keys: image
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: EnsureTyped
+    keys: image
+dataset:
+  _target_: Dataset
+  data: "@datalist"
+  transform: "@preprocessing"
+dataloader:
+  _target_: DataLoader
+  dataset: "@dataset"
+  batch_size: 1
+  shuffle: false
+  num_workers: 4
+inferer:
+  _target_: SlidingWindowInferer
+  roi_size:
+  - 96
+  - 96
+  - 96
+  sw_batch_size: 4
+  overlap: 0.625
+postprocessing:
+  _target_: Compose
+  transforms:
+  - _target_: Activationsd
+    keys: pred
+    softmax: true
+  - _target_: Invertd
+    keys: pred
+    transform: "@preprocessing"
+    orig_keys: image
+    meta_key_postfix: meta_dict
+    nearest_interp: false
+    to_tensor: true
+  - _target_: AsDiscreted
+    keys: pred
+    argmax: true
+  - _target_: SaveImaged
+    keys: pred
+    meta_keys: pred_meta_dict
+    output_dir: "@output_dir"
+handlers:
+- _target_: CheckpointLoader
+  load_path: "$@bundle_root + '/models/model.pt'"
+  load_dict:
+    model: "@network"
+- _target_: StatsHandler
+  iteration_log: false
+evaluator:
+  _target_: SupervisedEvaluator
+  device: "@device"
+  val_data_loader: "@dataloader"
+  network: "@network"
+  inferer: "@inferer"
+  postprocessing: "@postprocessing"
+  val_handlers: "@handlers"
+  amp: true
+evaluating:
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+- "$@evaluator.run()"

configs/logging.conf

@@ -0,0 +1,21 @@
+[loggers]
+keys=root
+
+[handlers]
+keys=consoleHandler
+
+[formatters]
+keys=fullFormatter
+
+[logger_root]
+level=INFO
+handlers=consoleHandler
+
+[handler_consoleHandler]
+class=StreamHandler
+level=INFO
+formatter=fullFormatter
+args=(sys.stdout,)
+
+[formatter_fullFormatter]
+format=%(asctime)s - %(name)s - %(levelname)s - %(message)s

configs/metadata.json

@@ -0,0 +1,80 @@
+{
+    "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
+    "version": "0.1.0",
+    "changelog": {
+        "0.1.0": "complete the model package",
+        "0.0.1": "initialize the model package structure"
+    },
+    "monai_version": "0.9.0",
+    "pytorch_version": "1.12.0",
+    "numpy_version": "1.21.2",
+    "optional_packages_version": {
+        "fire": "0.4.0",
+        "nibabel": "3.2.1",
+        "pytorch-ignite": "0.4.9"
+    },
+    "task": "Neural architecture search on pancreas CT segmentation",
+    "description": "Searched architectures for volumetric (3D) segmentation of the pancreas from CT image",
+    "authors": "MONAI team",
+    "copyright": "Copyright (c) MONAI Consortium",
+    "data_source": "Task07_Pancreas.tar from http://medicaldecathlon.com/",
+    "data_type": "nibabel",
+    "image_classes": "single channel data, intensity scaled to [0, 1]",
+    "label_classes": "single channel data, 1 is pancreas, 2 is pancreatic tumor, 0 is everything else",
+    "pred_classes": "3 channels OneHot data, channel 1 is pancreas, channel 2 is pancreatic tumor, channel 0 is background",
+    "eval_metrics": {
+        "mean_dice": 0.72
+    },
+    "intended_use": "This is an example, not to be used for diagnostic purposes",
+    "references": [
+        "He, Y., Yang, D., Roth, H., Zhao, C. and Xu, D., 2021. Dints: Differentiable neural network topology search for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 5841-5850)."
+    ],
+    "network_data_format": {
+        "inputs": {
+            "image": {
+                "type": "image",
+                "format": "hounsfield",
+                "modality": "CT",
+                "num_channels": 1,
+                "spatial_shape": [
+                    96,
+                    96,
+                    96
+                ],
+                "dtype": "float32",
+                "value_range": [
+                    0,
+                    1
+                ],
+                "is_patch_data": true,
+                "channel_def": {
+                    "0": "image"
+                }
+            }
+        },
+        "outputs": {
+            "pred": {
+                "type": "image",
+                "format": "segmentation",
+                "num_channels": 3,
+                "spatial_shape": [
+                    96,
+                    96,
+                    96
+                ],
+                "dtype": "float32",
+                "value_range": [
+                    0,
+                    1,
+                    2
+                ],
+                "is_patch_data": true,
+                "channel_def": {
+                    "0": "background",
+                    "1": "pancreas",
+                    "2": "pancreatic tumor"
+                }
+            }
+        }
+    }
+}

configs/multi_gpu_train.yaml

@@ -0,0 +1,52 @@
+---
+device: "$torch.device(f'cuda:{dist.get_rank()}')"
+network:
+  _target_: torch.nn.parallel.DistributedDataParallel
+  module: "$@network_def.to(@device)"
+  find_unused_parameters: true
+  device_ids:
+  - "@device"
+optimizer#lr: "$0.0125*dist.get_world_size()"
+lr_scheduler#step_size: "$80*dist.get_world_size()"
+train#handlers:
+  - _target_: LrScheduleHandler
+    lr_scheduler: "@lr_scheduler"
+    print_lr: true
+  - _target_: ValidationHandler
+    validator: "@validate#evaluator"
+    epoch_level: true
+    interval: "$10*dist.get_world_size()"
+  - _target_: StatsHandler
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+  - _target_: TensorBoardStatsHandler
+    log_dir: "@output_dir"
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+train#trainer#max_epochs: "$400*dist.get_world_size()"
+train#trainer#train_handlers: "$@train#handlers[: -2 if dist.get_rank() > 0 else None]"
+validate#evaluator#val_handlers: "$None if dist.get_rank() > 0 else @validate#handlers"
+training:
+- "$import torch.distributed as dist"
+- "$dist.init_process_group(backend='nccl')"
+- "$torch.cuda.set_device(@device)"
+- "$monai.utils.set_determinism(seed=123)"
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+- "$@train#trainer.run()"
+- "$dist.destroy_process_group()"
+train_data_partition: "$monai.data.partition_dataset(data=@train_datalist, num_partitions=dist.get_world_size(),
+  shuffle=True, even_divisible=True,)[dist.get_rank()]"
+train#dataset:
+  _target_: CacheDataset
+  data: "@train_data_partition"
+  transform: "@train#preprocessing"
+  cache_rate: 1
+  num_workers: 4
+val_data_partition: "$monai.data.partition_dataset(data=@val_datalist, num_partitions=dist.get_world_size(),
+  shuffle=False, even_divisible=False,)[dist.get_rank()]"
+validate#dataset:
+  _target_: CacheDataset
+  data: "@val_data_partition"
+  transform: "@validate#preprocessing"
+  cache_rate: 1
+  num_workers: 4

configs/search.yaml

@@ -0,0 +1,275 @@
+---
+imports:
+  - "$from scipy import ndimage"
+arch_ckpt_path: models
+amp: true
+data_file_base_dir: /workspace/data/msd/Task07_Pancreas
+data_list_file_path: configs/dataset_0.json
+determ: true
+input_channels: 1
+learning_rate: 0.025
+learning_rate_arch: 0.001
+learning_rate_milestones:
+- 0.4
+- 0.8
+num_images_per_batch: 1
+num_epochs: 1430
+num_epochs_per_validation: 100
+num_epochs_warmup: 715
+num_patches_per_image: 1
+num_sw_batch_size: 6
+output_classes: 3
+overlap_ratio: 0.625
+patch_size:
+- 96
+- 96
+- 96
+patch_size_valid:
+- 96
+- 96
+- 96
+ram_cost_factor: 0.8
+image_key: image
+label_key: label
+transform_train:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: EnsureChannelFirstd
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: Orientationd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    axcodes: RAS
+  - _target_: Spacingd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode:
+    - bilinear
+    - nearest
+    align_corners:
+    - true
+    - true
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: ScaleIntensityRanged
+    keys: "@image_key"
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$np.float16"
+    - "$np.uint8"
+  - _target_: CopyItemsd
+    keys: "@label_key"
+    times: 1
+    names:
+    - label4crop
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x, s=@output_classes: np.concatenate(tuple([ndimage.binary_dilation((x==_k).astype(x.dtype), iterations=48).astype(x.dtype) for _k in range(s)]), axis=0)"
+    overwrite: true
+  - _target_: EnsureTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: SpatialPadd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    - label4crop
+    spatial_size: "@patch_size"
+    mode:
+    - reflect
+    - constant
+    - constant
+  - _target_: RandCropByLabelClassesd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    label_key: label4crop
+    num_classes: "@output_classes"
+    ratios: "$[1,] * @output_classes"
+    spatial_size: "@patch_size"
+    num_samples: "@num_patches_per_image"
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x: 0"
+  - _target_: RandRotated
+    keys:
+    - "@image_key"
+    - "@label_key"
+    range_x: 0.3
+    range_y: 0.3
+    range_z: 0.3
+    mode:
+    - bilinear
+    - nearest
+    prob: 0.2
+  - _target_: RandZoomd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    min_zoom: 0.8
+    max_zoom: 1.2
+    mode:
+    - trilinear
+    - nearest
+    align_corners:
+    - null
+    - null
+    prob: 0.16
+  - _target_: RandGaussianSmoothd
+    keys: "@image_key"
+    sigma_x:
+    - 0.5
+    - 1.15
+    sigma_y:
+    - 0.5
+    - 1.15
+    sigma_z:
+    - 0.5
+    - 1.15
+    prob: 0.15
+  - _target_: RandScaleIntensityd
+    keys: "@image_key"
+    factors: 0.3
+    prob: 0.5
+  - _target_: RandShiftIntensityd
+    keys: "@image_key"
+    offsets: 0.1
+    prob: 0.5
+  - _target_: RandGaussianNoised
+    keys: "@image_key"
+    std: 0.01
+    prob: 0.15
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 0
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 1
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 2
+    prob: 0.5
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$torch.float32"
+    - "$torch.uint8"
+  - _target_: ToTensord
+    keys:
+    - "@image_key"
+    - "@label_key"
+transform_validation:
+  _target_: Compose
+  transforms:
+  - _target_: LoadImaged
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: EnsureChannelFirstd
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: Orientationd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    axcodes: RAS
+  - _target_: Spacingd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode:
+    - bilinear
+    - nearest
+    align_corners:
+    - true
+    - true
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: ScaleIntensityRanged
+    keys: "@image_key"
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$np.float16"
+    - "$np.uint8"
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$torch.float32"
+    - "$torch.uint8"
+  - _target_: ToTensord
+    keys:
+    - "@image_key"
+    - "@label_key"
+loss:
+  _target_: DiceCELoss
+  include_background: false
+  to_onehot_y: true
+  softmax: true
+  squared_pred: true
+  batch: true
+  smooth_nr: 0.00001
+  smooth_dr: 0.00001
+dints_space:
+  _target_: monai.networks.nets.TopologySearch
+  channel_mul: 0.5
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  device: "$torch.device('cuda')"
+network:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true

configs/train.yaml

@@ -0,0 +1,353 @@
+---
+imports:
+- "$import glob"
+- "$import json"
+- "$import os"
+- "$import ignite"
+- "$from scipy import ndimage"
+input_channels: 1
+output_classes: 3
+arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
+arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
+bundle_root: "/workspace/MONAI/model-zoo/models/pancreas_ct_dints_segmentation"
+ckpt_dir: "$@bundle_root + '/models'"
+output_dir: "$@bundle_root + '/eval'"
+dataset_dir: "/workspace/data/msd/Task07_Pancreas"
+data_list_file_path: "$@bundle_root + '/configs/dataset_0.json'"
+train_datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='training',
+  base_dir=@dataset_dir)"
+val_datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='validation',
+  base_dir=@dataset_dir)"
+device: "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
+dints_space:
+  _target_: monai.networks.nets.TopologyInstance
+  channel_mul: 1
+  num_blocks: 12
+  num_depths: 4
+  use_downsample: true
+  arch_code:
+  - "$@arch_ckpt['arch_code_a']"
+  - "$@arch_ckpt['arch_code_c']"
+  device: "$torch.device('cuda')"
+network_def:
+  _target_: monai.networks.nets.DiNTS
+  dints_space: "@dints_space"
+  in_channels: "@input_channels"
+  num_classes: "@output_classes"
+  use_downsample: true
+  node_a: "$@arch_ckpt['node_a']"
+network: "$@network_def.to(@device)"
+loss:
+  _target_: DiceCELoss
+  include_background: false
+  to_onehot_y: true
+  softmax: true
+  squared_pred: true
+  batch: true
+  smooth_nr: 1.0e-05
+  smooth_dr: 1.0e-05
+optimizer:
+  _target_: torch.optim.SGD
+  params: "$@network.parameters()"
+  momentum: 0.9
+  weight_decay: 4.0e-05
+  lr: 0.025
+lr_scheduler:
+  _target_: torch.optim.lr_scheduler.StepLR
+  optimizer: "@optimizer"
+  step_size: 80
+  gamma: 0.5
+image_key: image
+label_key: label
+train:
+  deterministic_transforms:
+  - _target_: LoadImaged
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: EnsureChannelFirstd
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: Orientationd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    axcodes: RAS
+  - _target_: Spacingd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    pixdim:
+    - 1
+    - 1
+    - 1
+    mode:
+    - bilinear
+    - nearest
+    align_corners:
+    - true
+    - true
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: ScaleIntensityRanged
+    keys: "@image_key"
+    a_min: -87
+    a_max: 199
+    b_min: 0
+    b_max: 1
+    clip: true
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$np.float16"
+    - "$np.uint8"
+  - _target_: CopyItemsd
+    keys: "@label_key"
+    times: 1
+    names:
+    - label4crop
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x, s=@output_classes: np.concatenate(tuple([ndimage.binary_dilation((x==_k).astype(x.dtype),
+      iterations=48).astype(x.dtype) for _k in range(s)]), axis=0)"
+    overwrite: true
+  - _target_: EnsureTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+  - _target_: CastToTyped
+    keys: "@image_key"
+    dtype: "$torch.float32"
+  - _target_: SpatialPadd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    - label4crop
+    spatial_size:
+    - 96
+    - 96
+    - 96
+    mode:
+    - reflect
+    - constant
+    - constant
+  random_transforms:
+  - _target_: RandCropByLabelClassesd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    label_key: label4crop
+    num_classes: "@output_classes"
+    ratios: "$[1,] * @output_classes"
+    spatial_size:
+    - 96
+    - 96
+    - 96
+    num_samples: 1
+  - _target_: Lambdad
+    keys: label4crop
+    func: "$lambda x: 0"
+  - _target_: RandRotated
+    keys:
+    - "@image_key"
+    - "@label_key"
+    range_x: 0.3
+    range_y: 0.3
+    range_z: 0.3
+    mode:
+    - bilinear
+    - nearest
+    prob: 0.2
+  - _target_: RandZoomd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    min_zoom: 0.8
+    max_zoom: 1.2
+    mode:
+    - trilinear
+    - nearest
+    align_corners:
+    - true
+    -
+    prob: 0.16
+  - _target_: RandGaussianSmoothd
+    keys: "@image_key"
+    sigma_x:
+    - 0.5
+    - 1.15
+    sigma_y:
+    - 0.5
+    - 1.15
+    sigma_z:
+    - 0.5
+    - 1.15
+    prob: 0.15
+  - _target_: RandScaleIntensityd
+    keys: "@image_key"
+    factors: 0.3
+    prob: 0.5
+  - _target_: RandShiftIntensityd
+    keys: "@image_key"
+    offsets: 0.1
+    prob: 0.5
+  - _target_: RandGaussianNoised
+    keys: "@image_key"
+    std: 0.01
+    prob: 0.15
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 0
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 1
+    prob: 0.5
+  - _target_: RandFlipd
+    keys:
+    - "@image_key"
+    - "@label_key"
+    spatial_axis: 2
+    prob: 0.5
+  - _target_: CastToTyped
+    keys:
+    - "@image_key"
+    - "@label_key"
+    dtype:
+    - "$torch.float32"
+    - "$torch.uint8"
+  - _target_: ToTensord
+    keys:
+    - "@image_key"
+    - "@label_key"
+  preprocessing:
+    _target_: Compose
+    transforms: "$@train#deterministic_transforms + @train#random_transforms"
+  dataset:
+    _target_: CacheDataset
+    data: "@train_datalist"
+    transform: "@train#preprocessing"
+    cache_rate: 0.125
+    num_workers: 4
+  dataloader:
+    _target_: DataLoader
+    dataset: "@train#dataset"
+    batch_size: 2
+    shuffle: true
+    num_workers: 4
+  inferer:
+    _target_: SimpleInferer
+  postprocessing:
+    _target_: Compose
+    transforms:
+    - _target_: Activationsd
+      keys: pred
+      softmax: true
+    - _target_: AsDiscreted
+      keys:
+      - pred
+      - label
+      argmax:
+      - true
+      - false
+      to_onehot: "@output_classes"
+  handlers:
+  - _target_: LrScheduleHandler
+    lr_scheduler: "@lr_scheduler"
+    print_lr: true
+  - _target_: ValidationHandler
+    validator: "@validate#evaluator"
+    epoch_level: true
+    interval: 10
+  - _target_: StatsHandler
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+  - _target_: TensorBoardStatsHandler
+    log_dir: "@output_dir"
+    tag_name: train_loss
+    output_transform: "$monai.handlers.from_engine(['loss'], first=True)"
+  key_metric:
+    train_accuracy:
+      _target_: ignite.metrics.Accuracy
+      output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+  trainer:
+    _target_: SupervisedTrainer
+    max_epochs: 400
+    device: "@device"
+    train_data_loader: "@train#dataloader"
+    network: "@network"
+    loss_function: "@loss"
+    optimizer: "@optimizer"
+    inferer: "@train#inferer"
+    postprocessing: "@train#postprocessing"
+    key_train_metric: "@train#key_metric"
+    train_handlers: "@train#handlers"
+    amp: true
+validate:
+  preprocessing:
+    _target_: Compose
+    transforms: "%train#deterministic_transforms"
+  dataset:
+    _target_: CacheDataset
+    data: "@val_datalist"
+    transform: "@validate#preprocessing"
+    cache_rate: 0.125
+  dataloader:
+    _target_: DataLoader
+    dataset: "@validate#dataset"
+    batch_size: 1
+    shuffle: false
+    num_workers: 4
+  inferer:
+    _target_: SlidingWindowInferer
+    roi_size:
+    - 96
+    - 96
+    - 96
+    sw_batch_size: 6
+    overlap: 0.625
+  postprocessing: "%train#postprocessing"
+  handlers:
+  - _target_: StatsHandler
+    iteration_log: false
+  - _target_: TensorBoardStatsHandler
+    log_dir: "@output_dir"
+    iteration_log: false
+  - _target_: CheckpointSaver
+    save_dir: "@ckpt_dir"
+    save_dict:
+      model: "@network"
+    save_key_metric: true
+    key_metric_filename: model.pt
+  key_metric:
+    val_mean_dice:
+      _target_: MeanDice
+      include_background: false
+      output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+  additional_metrics:
+    val_accuracy:
+      _target_: ignite.metrics.Accuracy
+      output_transform: "$monai.handlers.from_engine(['pred', 'label'])"
+  evaluator:
+    _target_: SupervisedEvaluator
+    device: "@device"
+    val_data_loader: "@validate#dataloader"
+    network: "@network"
+    inferer: "@validate#inferer"
+    postprocessing: "@validate#postprocessing"
+    key_val_metric: "@validate#key_metric"
+    additional_metrics: "@validate#additional_metrics"
+    val_handlers: "@validate#handlers"
+    amp: true
+training:
+- "$monai.utils.set_determinism(seed=123)"
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"
+- "$@train#trainer.run()"

docs/README.md

@@ -0,0 +1,72 @@
+# Description
+A neural architecture search algorithm for volumetric (3D) segmentation of the pancreas and pancreatic tumor from CT image.
+
+# Model Overview
+This model is trained using the state-of-the-art algorithm [1] of the "Medical Segmentation Decathlon Challenge 2018" with 196 training images, 56 validation images, and 28 testing images.
+
+## Data
+The training dataset is Task07_Pancreas.tar from http://medicaldecathlon.com/. And the data list/split can be created with the script `scripts/prepare_datalist.py`.
+
+## Training configuration
+The training was performed with at least 16GB-memory GPUs.
+
+Actual Model Input: 96 x 96 x 96
+
+## Input and output formats
+Input: 1 channel CT image
+
+Output: 3 channels: Label 2: pancreatic tumor; Label 1: pancreas; Label 0: everything else
+
+## Scores
+This model achieves the following Dice score on the validation data (our own split from the training dataset):
+
+Mean Dice = 0.72
+
+## commands example
+Create data split (.json file):
+
+```
+python scripts/prepare_datalist.py --path /path-to-Task07_Pancreas/ --output configs/dataset_0.json
+```
+
+Execute model searching:
+
+```
+python -m scripts.search run --config_file configs/search.yaml
+```
+
+Execute multi-GPU model searching (recommended):
+
+```
+torchrun --nnodes=1 --nproc_per_node=8 -m scripts.search run --config_file configs/search.yaml
+```
+
+Execute training:
+
+```
+python -m monai.bundle run training --meta_file configs/metadata.json --config_file configs/train.yaml --logging_file configs/logging.conf
+```
+
+Override the `train` config to execute multi-GPU training:
+
+```
+torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run training --meta_file configs/metadata.json --config_file "['configs/train.yaml','configs/multi_gpu_train.yaml']" --logging_file configs/logging.conf
+```
+
+Override the `train` config to execute evaluation with the trained model:
+
+```
+python -m monai.bundle run evaluating --meta_file configs/metadata.json --config_file "['configs/train.yaml','configs/evaluate.yaml']" --logging_file configs/logging.conf
+```
+
+Execute inference:
+
+```
+python -m monai.bundle run evaluating --meta_file configs/metadata.json --config_file configs/inference.yaml --logging_file configs/logging.conf
+```
+
+# Disclaimer
+This is an example, not to be used for diagnostic purposes.
+
+# References
+[1] He, Y., Yang, D., Roth, H., Zhao, C. and Xu, D., 2021. Dints: Differentiable neural network topology search for 3d medical image segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 5841-5850).

docs/license.txt

@@ -0,0 +1,6 @@
+Third Party Licenses
+-----------------------------------------------------------------------
+
+/*********************************************************************/
+i. Medical Segmentation Decathlon
+   http://medicaldecathlon.com/

models/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:975201057eb16225a1abfd047cb9b293f6d481dc604468d512710f3543f29066
+size 616210421

models/model.ts

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:526d2bdb4d88f6f55f2d88eb0c79deeeb90b3ced182269e33cdc5da6e46ea5fb
+size 616338455

models/search_code_18590.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01e361e9843e2f4e5ff1599da0abac77013ea38cab8fdd6c9286bb6572c9a32d
+size 4335

scripts/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

scripts/prepare_datalist.py

@@ -0,0 +1,59 @@
+import argparse
+import glob
+import json
+import os
+
+import monai
+from sklearn.model_selection import train_test_split
+
+
+def produce_sample_dict(line: str):
+    return {"label": line, "image": line.replace("labelsTr", "imagesTr")}
+
+
+def produce_datalist(dataset_dir: str):
+    """
+    This function is used to split the dataset.
+    It will produce 200 samples for training, and the other samples are divided equally
+    into val and test sets.
+    """
+
+    samples = sorted(glob.glob(os.path.join(dataset_dir, "labelsTr", "*"), recursive=True))
+    samples = [_item.replace(os.path.join(dataset_dir, "labelsTr"), "labelsTr") for _item in samples]
+    datalist = []
+    for line in samples:
+        datalist.append(produce_sample_dict(line))
+    train_list, other_list = train_test_split(datalist, train_size=196)
+    val_list, test_list = train_test_split(other_list, train_size=0.66)
+
+    return {"training": train_list, "validation": val_list, "testing": test_list}
+
+
+def main(args):
+    """
+    split the dataset and output the data list into a json file.
+    """
+    data_file_base_dir = args.path
+    output_json = args.output
+    # produce deterministic data splits
+    monai.utils.set_determinism(seed=123)
+    datalist = produce_datalist(dataset_dir=data_file_base_dir)
+    with open(output_json, "w") as f:
+        json.dump(datalist, f, ensure_ascii=True, indent=4)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "--path",
+        type=str,
+        default="/workspace/data/msd/Task07_Pancreas",
+        help="root path of MSD Task07_Pancreas dataset.",
+    )
+    parser.add_argument(
+        "--output", type=str, default="dataset_0.json", help="relative path of output datalist json file."
+    )
+    args = parser.parse_args()
+
+    main(args)

scripts/search.py

@@ -0,0 +1,520 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import os
+import random
+import sys
+import time
+from datetime import datetime
+from typing import Sequence, Union
+
+import monai
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import yaml
+from monai import transforms
+from monai.bundle import ConfigParser
+from monai.data import ThreadDataLoader, partition_dataset
+from monai.inferers import sliding_window_inference
+from monai.metrics import compute_meandice
+from monai.utils import set_determinism
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.tensorboard import SummaryWriter
+
+
+def run(config_file: Union[str, Sequence[str]]):
+    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+    parser = ConfigParser()
+    parser.read_config(config_file)
+
+    arch_ckpt_path = parser["arch_ckpt_path"]
+    amp = parser["amp"]
+    data_file_base_dir = parser["data_file_base_dir"]
+    data_list_file_path = parser["data_list_file_path"]
+    determ = parser["determ"]
+    learning_rate = parser["learning_rate"]
+    learning_rate_arch = parser["learning_rate_arch"]
+    learning_rate_milestones = np.array(parser["learning_rate_milestones"])
+    num_images_per_batch = parser["num_images_per_batch"]
+    num_epochs = parser["num_epochs"]  # around 20k iterations
+    num_epochs_per_validation = parser["num_epochs_per_validation"]
+    num_epochs_warmup = parser["num_epochs_warmup"]
+    num_sw_batch_size = parser["num_sw_batch_size"]
+    output_classes = parser["output_classes"]
+    overlap_ratio = parser["overlap_ratio"]
+    patch_size_valid = parser["patch_size_valid"]
+    ram_cost_factor = parser["ram_cost_factor"]
+    print("[info] GPU RAM cost factor:", ram_cost_factor)
+
+    train_transforms = parser.get_parsed_content("transform_train")
+    val_transforms = parser.get_parsed_content("transform_validation")
+
+    # deterministic training
+    if determ:
+        set_determinism(seed=0)
+
+    print("[info] number of GPUs:", torch.cuda.device_count())
+    if torch.cuda.device_count() > 1:
+        # initialize the distributed training process, every GPU runs in a process
+        dist.init_process_group(backend="nccl", init_method="env://")
+        world_size = dist.get_world_size()
+    else:
+        world_size = 1
+    print("[info] world_size:", world_size)
+
+    with open(data_list_file_path, "r") as f:
+        json_data = json.load(f)
+
+    list_train = json_data["training"]
+    list_valid = json_data["validation"]
+
+    # training data
+    files = []
+    for _i in range(len(list_train)):
+        str_img = os.path.join(data_file_base_dir, list_train[_i]["image"])
+        str_seg = os.path.join(data_file_base_dir, list_train[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    train_files = files
+
+    random.shuffle(train_files)
+
+    train_files_w = train_files[: len(train_files) // 2]
+    if torch.cuda.device_count() > 1:
+        train_files_w = partition_dataset(
+            data=train_files_w, shuffle=True, num_partitions=world_size, even_divisible=True
+        )[dist.get_rank()]
+    print("train_files_w:", len(train_files_w))
+
+    train_files_a = train_files[len(train_files) // 2 :]
+    if torch.cuda.device_count() > 1:
+        train_files_a = partition_dataset(
+            data=train_files_a, shuffle=True, num_partitions=world_size, even_divisible=True
+        )[dist.get_rank()]
+    print("train_files_a:", len(train_files_a))
+
+    # validation data
+    files = []
+    for _i in range(len(list_valid)):
+        str_img = os.path.join(data_file_base_dir, list_valid[_i]["image"])
+        str_seg = os.path.join(data_file_base_dir, list_valid[_i]["label"])
+
+        if (not os.path.exists(str_img)) or (not os.path.exists(str_seg)):
+            continue
+
+        files.append({"image": str_img, "label": str_seg})
+    val_files = files
+
+    if torch.cuda.device_count() > 1:
+        val_files = partition_dataset(data=val_files, shuffle=False, num_partitions=world_size, even_divisible=False)[
+            dist.get_rank()
+        ]
+    print("val_files:", len(val_files))
+
+    # network architecture
+    if torch.cuda.device_count() > 1:
+        device = torch.device(f"cuda:{dist.get_rank()}")
+    else:
+        device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
+
+    if torch.cuda.device_count() > 1:
+        train_ds_a = monai.data.CacheDataset(
+            data=train_files_a, transform=train_transforms, cache_rate=1.0, num_workers=8
+        )
+        train_ds_w = monai.data.CacheDataset(
+            data=train_files_w, transform=train_transforms, cache_rate=1.0, num_workers=8
+        )
+        val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=2)
+    else:
+        train_ds_a = monai.data.CacheDataset(
+            data=train_files_a, transform=train_transforms, cache_rate=0.125, num_workers=8
+        )
+        train_ds_w = monai.data.CacheDataset(
+            data=train_files_w, transform=train_transforms, cache_rate=0.125, num_workers=8
+        )
+        val_ds = monai.data.CacheDataset(data=val_files, transform=val_transforms, cache_rate=0.125, num_workers=2)
+
+    train_loader_a = ThreadDataLoader(train_ds_a, num_workers=6, batch_size=num_images_per_batch, shuffle=True)
+    train_loader_w = ThreadDataLoader(train_ds_w, num_workers=6, batch_size=num_images_per_batch, shuffle=True)
+    val_loader = ThreadDataLoader(val_ds, num_workers=0, batch_size=1, shuffle=False)
+
+    model = parser.get_parsed_content("network")
+    dints_space = parser.get_parsed_content("dints_space")
+
+    model = model.to(device)
+    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    post_pred = transforms.Compose(
+        [transforms.EnsureType(), transforms.AsDiscrete(argmax=True, to_onehot=output_classes)]
+    )
+    post_label = transforms.Compose([transforms.EnsureType(), transforms.AsDiscrete(to_onehot=output_classes)])
+
+    # loss function
+    loss_func = parser.get_parsed_content("loss")
+
+    # optimizer
+    optimizer = torch.optim.SGD(
+        model.weight_parameters(), lr=learning_rate * world_size, momentum=0.9, weight_decay=0.00004
+    )
+    arch_optimizer_a = torch.optim.Adam(
+        [dints_space.log_alpha_a], lr=learning_rate_arch * world_size, betas=(0.5, 0.999), weight_decay=0.0
+    )
+    arch_optimizer_c = torch.optim.Adam(
+        [dints_space.log_alpha_c], lr=learning_rate_arch * world_size, betas=(0.5, 0.999), weight_decay=0.0
+    )
+
+    if torch.cuda.device_count() > 1:
+        model = DistributedDataParallel(model, device_ids=[device], find_unused_parameters=True)
+
+    # amp
+    if amp:
+        from torch.cuda.amp import GradScaler, autocast
+
+        scaler = GradScaler()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("[info] amp enabled")
+
+    # start a typical PyTorch training
+    val_interval = num_epochs_per_validation
+    best_metric = -1
+    best_metric_epoch = -1
+    idx_iter = 0
+
+    if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+        writer = SummaryWriter(log_dir=os.path.join(arch_ckpt_path, "Events"))
+
+        with open(os.path.join(arch_ckpt_path, "accuracy_history.csv"), "a") as f:
+            f.write("epoch\tmetric\tloss\tlr\ttime\titer\n")
+
+    dataloader_a_iterator = iter(train_loader_a)
+
+    start_time = time.time()
+    for epoch in range(num_epochs):
+        decay = 0.5 ** np.sum(
+            [(epoch - num_epochs_warmup) / (num_epochs - num_epochs_warmup) > learning_rate_milestones]
+        )
+        lr = learning_rate * decay * world_size
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            print("-" * 10)
+            print(f"epoch {epoch + 1}/{num_epochs}")
+            print("learning rate is set to {}".format(lr))
+
+        model.train()
+        epoch_loss = 0
+        loss_torch = torch.zeros(2, dtype=torch.float, device=device)
+        epoch_loss_arch = 0
+        loss_torch_arch = torch.zeros(2, dtype=torch.float, device=device)
+        step = 0
+
+        for batch_data in train_loader_w:
+            step += 1
+            inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
+            if world_size == 1:
+                for _ in model.weight_parameters():
+                    _.requires_grad = True
+            else:
+                for _ in model.module.weight_parameters():
+                    _.requires_grad = True
+            dints_space.log_alpha_a.requires_grad = False
+            dints_space.log_alpha_c.requires_grad = False
+
+            optimizer.zero_grad()
+
+            if amp:
+                with autocast():
+                    outputs = model(inputs)
+                    if output_classes == 2:
+                        loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                    else:
+                        loss = loss_func(outputs, labels)
+
+                scaler.scale(loss).backward()
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                outputs = model(inputs)
+                if output_classes == 2:
+                    loss = loss_func(torch.flip(outputs, dims=[1]), 1 - labels)
+                else:
+                    loss = loss_func(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+            epoch_loss += loss.item()
+            loss_torch[0] += loss.item()
+            loss_torch[1] += 1.0
+            epoch_len = len(train_loader_w)
+            idx_iter += 1
+
+            if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+                print("[{0}] ".format(str(datetime.now())[:19]) + f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
+                writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
+
+            if epoch < num_epochs_warmup:
+                continue
+
+            try:
+                sample_a = next(dataloader_a_iterator)
+            except StopIteration:
+                dataloader_a_iterator = iter(train_loader_a)
+                sample_a = next(dataloader_a_iterator)
+            inputs_search, labels_search = (sample_a["image"].to(device), sample_a["label"].to(device))
+            if world_size == 1:
+                for _ in model.weight_parameters():
+                    _.requires_grad = False
+            else:
+                for _ in model.module.weight_parameters():
+                    _.requires_grad = False
+            dints_space.log_alpha_a.requires_grad = True
+            dints_space.log_alpha_c.requires_grad = True
+
+            # linear increase topology and RAM loss
+            entropy_alpha_c = torch.tensor(0.0).to(device)
+            entropy_alpha_a = torch.tensor(0.0).to(device)
+            ram_cost_full = torch.tensor(0.0).to(device)
+            ram_cost_usage = torch.tensor(0.0).to(device)
+            ram_cost_loss = torch.tensor(0.0).to(device)
+            topology_loss = torch.tensor(0.0).to(device)
+
+            probs_a, arch_code_prob_a = dints_space.get_prob_a(child=True)
+            entropy_alpha_a = -((probs_a) * torch.log(probs_a + 1e-5)).mean()
+            entropy_alpha_c = -(
+                F.softmax(dints_space.log_alpha_c, dim=-1) * F.log_softmax(dints_space.log_alpha_c, dim=-1)
+            ).mean()
+            topology_loss = dints_space.get_topology_entropy(probs_a)
+
+            ram_cost_full = dints_space.get_ram_cost_usage(inputs.shape, full=True)
+            ram_cost_usage = dints_space.get_ram_cost_usage(inputs.shape)
+            ram_cost_loss = torch.abs(ram_cost_factor - ram_cost_usage / ram_cost_full)
+
+            arch_optimizer_a.zero_grad()
+            arch_optimizer_c.zero_grad()
+
+            combination_weights = (epoch - num_epochs_warmup) / (num_epochs - num_epochs_warmup)
+
+            if amp:
+                with autocast():
+                    outputs_search = model(inputs_search)
+                    if output_classes == 2:
+                        loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                    else:
+                        loss = loss_func(outputs_search, labels_search)
+
+                    loss += combination_weights * (
+                        (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                    )
+
+                scaler.scale(loss).backward()
+                scaler.step(arch_optimizer_a)
+                scaler.step(arch_optimizer_c)
+                scaler.update()
+            else:
+                outputs_search = model(inputs_search)
+                if output_classes == 2:
+                    loss = loss_func(torch.flip(outputs_search, dims=[1]), 1 - labels_search)
+                else:
+                    loss = loss_func(outputs_search, labels_search)
+
+                loss += 1.0 * (
+                    combination_weights * (entropy_alpha_a + entropy_alpha_c) + ram_cost_loss + 0.001 * topology_loss
+                )
+
+                loss.backward()
+                arch_optimizer_a.step()
+                arch_optimizer_c.step()
+
+            epoch_loss_arch += loss.item()
+            loss_torch_arch[0] += loss.item()
+            loss_torch_arch[1] += 1.0
+
+            if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+                print(
+                    "[{0}] ".format(str(datetime.now())[:19])
+                    + f"{step}/{epoch_len}, train_loss_arch: {loss.item():.4f}"
+                )
+                writer.add_scalar("train_loss_arch", loss.item(), epoch_len * epoch + step)
+
+        # synchronizes all processes and reduce results
+        if torch.cuda.device_count() > 1:
+            dist.barrier()
+            dist.all_reduce(loss_torch, op=torch.distributed.ReduceOp.SUM)
+
+        loss_torch = loss_torch.tolist()
+        loss_torch_arch = loss_torch_arch.tolist()
+        if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+            loss_torch_epoch = loss_torch[0] / loss_torch[1]
+            print(
+                f"epoch {epoch + 1} average loss: {loss_torch_epoch:.4f}, "
+                f"best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
+            )
+
+            if epoch >= num_epochs_warmup:
+                loss_torch_arch_epoch = loss_torch_arch[0] / loss_torch_arch[1]
+                print(
+                    f"epoch {epoch + 1} average arch loss: {loss_torch_arch_epoch:.4f}, "
+                    f"best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}"
+                )
+
+        if (epoch + 1) % val_interval == 0 or (epoch + 1) == num_epochs:
+            torch.cuda.empty_cache()
+            model.eval()
+            with torch.no_grad():
+                metric = torch.zeros((output_classes - 1) * 2, dtype=torch.float, device=device)
+                metric_sum = 0.0
+                metric_count = 0
+                metric_mat = []
+                val_images = None
+                val_labels = None
+                val_outputs = None
+
+                _index = 0
+                for val_data in val_loader:
+                    val_images = val_data["image"].to(device)
+                    val_labels = val_data["label"].to(device)
+
+                    roi_size = patch_size_valid
+                    sw_batch_size = num_sw_batch_size
+
+                    if amp:
+                        with torch.cuda.amp.autocast():
+                            pred = sliding_window_inference(
+                                val_images,
+                                roi_size,
+                                sw_batch_size,
+                                lambda x: model(x),
+                                mode="gaussian",
+                                overlap=overlap_ratio,
+                            )
+                    else:
+                        pred = sliding_window_inference(
+                            val_images,
+                            roi_size,
+                            sw_batch_size,
+                            lambda x: model(x),
+                            mode="gaussian",
+                            overlap=overlap_ratio,
+                        )
+                    val_outputs = pred
+
+                    val_outputs = post_pred(val_outputs[0, ...])
+                    val_outputs = val_outputs[None, ...]
+                    val_labels = post_label(val_labels[0, ...])
+                    val_labels = val_labels[None, ...]
+
+                    value = compute_meandice(y_pred=val_outputs, y=val_labels, include_background=False)
+
+                    print(_index + 1, "/", len(val_loader), value)
+
+                    metric_count += len(value)
+                    metric_sum += value.sum().item()
+                    metric_vals = value.cpu().numpy()
+                    if len(metric_mat) == 0:
+                        metric_mat = metric_vals
+                    else:
+                        metric_mat = np.concatenate((metric_mat, metric_vals), axis=0)
+
+                    for _c in range(output_classes - 1):
+                        val0 = torch.nan_to_num(value[0, _c], nan=0.0)
+                        val1 = 1.0 - torch.isnan(value[0, 0]).float()
+                        metric[2 * _c] += val0 * val1
+                        metric[2 * _c + 1] += val1
+
+                    _index += 1
+
+                # synchronizes all processes and reduce results
+                if torch.cuda.device_count() > 1:
+                    dist.barrier()
+                    dist.all_reduce(metric, op=torch.distributed.ReduceOp.SUM)
+
+                metric = metric.tolist()
+                if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+                    for _c in range(output_classes - 1):
+                        print("evaluation metric - class {0:d}:".format(_c + 1), metric[2 * _c] / metric[2 * _c + 1])
+                    avg_metric = 0
+                    for _c in range(output_classes - 1):
+                        avg_metric += metric[2 * _c] / metric[2 * _c + 1]
+                    avg_metric = avg_metric / float(output_classes - 1)
+                    print("avg_metric", avg_metric)
+
+                    if avg_metric > best_metric:
+                        best_metric = avg_metric
+                        best_metric_epoch = epoch + 1
+                        best_metric_iterations = idx_iter
+
+                    (node_a_d, arch_code_a_d, arch_code_c_d, arch_code_a_max_d) = dints_space.decode()
+                    torch.save(
+                        {
+                            "node_a": node_a_d,
+                            "arch_code_a": arch_code_a_d,
+                            "arch_code_a_max": arch_code_a_max_d,
+                            "arch_code_c": arch_code_c_d,
+                            "iter_num": idx_iter,
+                            "epochs": epoch + 1,
+                            "best_dsc": best_metric,
+                            "best_path": best_metric_iterations,
+                        },
+                        os.path.join(arch_ckpt_path, "search_code_" + str(idx_iter) + ".pt"),
+                    )
+                    print("saved new best metric model")
+
+                    dict_file = {}
+                    dict_file["best_avg_dice_score"] = float(best_metric)
+                    dict_file["best_avg_dice_score_epoch"] = int(best_metric_epoch)
+                    dict_file["best_avg_dice_score_iteration"] = int(idx_iter)
+                    with open(os.path.join(arch_ckpt_path, "progress.yaml"), "w") as out_file:
+                        _ = yaml.dump(dict_file, stream=out_file)
+
+                    print(
+                        "current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format(
+                            epoch + 1, avg_metric, best_metric, best_metric_epoch
+                        )
+                    )
+
+                    current_time = time.time()
+                    elapsed_time = (current_time - start_time) / 60.0
+                    with open(os.path.join(arch_ckpt_path, "accuracy_history.csv"), "a") as f:
+                        f.write(
+                            "{0:d}\t{1:.5f}\t{2:.5f}\t{3:.5f}\t{4:.1f}\t{5:d}\n".format(
+                                epoch + 1, avg_metric, loss_torch_epoch, lr, elapsed_time, idx_iter
+                            )
+                        )
+
+                if torch.cuda.device_count() > 1:
+                    dist.barrier()
+
+            torch.cuda.empty_cache()
+
+    print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
+
+    if torch.cuda.device_count() == 1 or dist.get_rank() == 0:
+        writer.close()
+
+    if torch.cuda.device_count() > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    from monai.utils import optional_import
+
+    fire, _ = optional_import("fire")
+    fire.Fire()