katielink's picture
Add training pipeline for fine-tuning models, support MONAI Label active learning
560f266
{
"imports": [
"$import glob",
"$import os",
"$import ignite"
],
"bundle_root": "/models/renalStructures_UNEST_segmentation",
"ckpt_dir": "$@bundle_root + '/models'",
"output_dir": "$@bundle_root + '/eval'",
"dataset_dir": "$@bundle_root + './dataset'",
"images": "$list(sorted(glob.glob(@dataset_dir + '/imagesTr/*.nii.gz')))",
"labels": "$list(sorted(glob.glob(@dataset_dir + '/labelsTr/*.nii.gz')))",
"val_interval": 5,
"device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
"network_def": {
"_target_": "scripts.networks.unest.UNesT",
"in_channels": 1,
"out_channels": 4
},
"network": "$@network_def.to(@device)",
"loss": {
"_target_": "DiceCELoss",
"to_onehot_y": true,
"softmax": true,
"squared_pred": true,
"batch": true
},
"optimizer": {
"_target_": "torch.optim.Adam",
"params": "$@network.parameters()",
"lr": 0.0002
},
"train": {
"deterministic_transforms": [
{
"_target_": "LoadImaged",
"keys": [
"image",
"label"
]
},
{
"_target_": "EnsureChannelFirstd",
"keys": [
"image",
"label"
]
},
{
"_target_": "Orientationd",
"keys": [
"image",
"label"
],
"axcodes": "RAS"
},
{
"_target_": "Spacingd",
"keys": [
"image",
"label"
],
"pixdim": [
1.0,
1.0,
1.0
],
"mode": [
"bilinear",
"nearest"
]
},
{
"_target_": "ScaleIntensityRanged",
"keys": "image",
"a_min": -175,
"a_max": 250,
"b_min": 0.0,
"b_max": 1.0,
"clip": true
},
{
"_target_": "EnsureTyped",
"keys": [
"image",
"label"
]
}
],
"random_transforms": [
{
"_target_": "RandCropByPosNegLabeld",
"keys": [
"image",
"label"
],
"label_key": "label",
"spatial_size": [
96,
96,
96
],
"pos": 1,
"neg": 1,
"num_samples": 4,
"image_key": "image",
"image_threshold": 0
},
{
"_target_": "RandFlipd",
"keys": [
"image",
"label"
],
"spatial_axis": [
0
],
"prob": 0.1
},
{
"_target_": "RandFlipd",
"keys": [
"image",
"label"
],
"spatial_axis": [
1
],
"prob": 0.1
},
{
"_target_": "RandFlipd",
"keys": [
"image",
"label"
],
"spatial_axis": [
2
],
"prob": 0.1
},
{
"_target_": "RandRotate90d",
"keys": [
"image",
"label"
],
"max_k": 3,
"prob": 0.1
},
{
"_target_": "RandShiftIntensityd",
"keys": "image",
"offsets": 0.1,
"prob": 0.5
}
],
"preprocessing": {
"_target_": "Compose",
"transforms": "$@train#deterministic_transforms + @train#random_transforms"
},
"dataset": {
"_target_": "CacheDataset",
"data": "$[{'image': i, 'label': l} for i, l in zip(@images[:-9], @labels[:-9])]",
"transform": "@train#preprocessing",
"cache_rate": 1.0,
"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": 4
}
]
},
"handlers": [
{
"_target_": "ValidationHandler",
"validator": "@validate#evaluator",
"epoch_level": true,
"interval": "@val_interval"
},
{
"_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": 1000,
"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": "$[{'image': i, 'label': l} for i, l in zip(@images[-9:], @labels[-9:])]",
"transform": "@validate#preprocessing",
"cache_rate": 1.0
},
"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": 4,
"overlap": 0.5
},
"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()"
]
}