| { |
| "imports": [ |
| "$import functools", |
| "$import glob", |
| "$import scripts" |
| ], |
| "bundle_root": ".", |
| "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", |
| "ckpt_dir": "$@bundle_root + '/models'", |
| "tf_dir": "$@bundle_root + '/eval'", |
| "dataset_dir": "/workspace/data/medical", |
| "pretrained": false, |
| "perceptual_loss_model_weights_path": null, |
| "train_batch_size": 2, |
| "lr": 1e-05, |
| "train_patch_size": [ |
| 112, |
| 128, |
| 80 |
| ], |
| "channel": 0, |
| "spacing": [ |
| 1.1, |
| 1.1, |
| 1.1 |
| ], |
| "spatial_dims": 3, |
| "image_channels": 1, |
| "latent_channels": 8, |
| "discriminator_def": { |
| "_target_": "monai.networks.nets.patchgan_discriminator.PatchDiscriminator", |
| "spatial_dims": "@spatial_dims", |
| "num_layers_d": 3, |
| "channels": 32, |
| "in_channels": 1, |
| "out_channels": 1, |
| "norm": "INSTANCE" |
| }, |
| "autoencoder_def": { |
| "_target_": "monai.networks.nets.autoencoderkl.AutoencoderKL", |
| "spatial_dims": "@spatial_dims", |
| "in_channels": "@image_channels", |
| "out_channels": "@image_channels", |
| "latent_channels": "@latent_channels", |
| "channels": [ |
| 64, |
| 128, |
| 256 |
| ], |
| "num_res_blocks": 2, |
| "norm_num_groups": 32, |
| "norm_eps": 1e-06, |
| "attention_levels": [ |
| false, |
| false, |
| false |
| ], |
| "with_encoder_nonlocal_attn": false, |
| "with_decoder_nonlocal_attn": false, |
| "include_fc": false |
| }, |
| "perceptual_loss_def": { |
| "_target_": "monai.losses.perceptual.PerceptualLoss", |
| "spatial_dims": "@spatial_dims", |
| "network_type": "resnet50", |
| "is_fake_3d": true, |
| "fake_3d_ratio": 0.2, |
| "pretrained": "@pretrained", |
| "pretrained_path": "@perceptual_loss_model_weights_path", |
| "pretrained_state_dict_key": "state_dict" |
| }, |
| "dnetwork": "$@discriminator_def.to(@device)", |
| "gnetwork": "$@autoencoder_def.to(@device)", |
| "loss_perceptual": "$@perceptual_loss_def.to(@device)", |
| "doptimizer": { |
| "_target_": "torch.optim.Adam", |
| "params": "$@dnetwork.parameters()", |
| "lr": "@lr" |
| }, |
| "goptimizer": { |
| "_target_": "torch.optim.Adam", |
| "params": "$@gnetwork.parameters()", |
| "lr": "@lr" |
| }, |
| "preprocessing_transforms": [ |
| { |
| "_target_": "LoadImaged", |
| "keys": "image" |
| }, |
| { |
| "_target_": "EnsureChannelFirstd", |
| "keys": "image" |
| }, |
| { |
| "_target_": "Lambdad", |
| "keys": "image", |
| "func": "$lambda x: x[@channel, :, :, :]" |
| }, |
| { |
| "_target_": "EnsureChannelFirstd", |
| "keys": "image", |
| "channel_dim": "no_channel" |
| }, |
| { |
| "_target_": "EnsureTyped", |
| "keys": "image" |
| }, |
| { |
| "_target_": "Orientationd", |
| "keys": "image", |
| "axcodes": "RAS" |
| }, |
| { |
| "_target_": "Spacingd", |
| "keys": "image", |
| "pixdim": "@spacing", |
| "mode": "bilinear" |
| } |
| ], |
| "final_transforms": [ |
| { |
| "_target_": "ScaleIntensityRangePercentilesd", |
| "keys": "image", |
| "lower": 0, |
| "upper": 99.5, |
| "b_min": 0, |
| "b_max": 1 |
| } |
| ], |
| "train": { |
| "crop_transforms": [ |
| { |
| "_target_": "RandSpatialCropd", |
| "keys": "image", |
| "roi_size": "@train_patch_size", |
| "random_size": false |
| } |
| ], |
| "preprocessing": { |
| "_target_": "Compose", |
| "transforms": "$@preprocessing_transforms + @train#crop_transforms + @final_transforms" |
| }, |
| "dataset": { |
| "_target_": "monai.apps.DecathlonDataset", |
| "root_dir": "@dataset_dir", |
| "task": "Task01_BrainTumour", |
| "section": "training", |
| "cache_rate": 1.0, |
| "num_workers": 8, |
| "download": false, |
| "transform": "@train#preprocessing" |
| }, |
| "dataloader": { |
| "_target_": "DataLoader", |
| "dataset": "@train#dataset", |
| "batch_size": "@train_batch_size", |
| "shuffle": true, |
| "num_workers": 0 |
| }, |
| "handlers": [ |
| { |
| "_target_": "CheckpointSaver", |
| "save_dir": "@ckpt_dir", |
| "save_dict": { |
| "model": "@gnetwork" |
| }, |
| "save_interval": 0, |
| "save_final": true, |
| "epoch_level": true, |
| "final_filename": "model_autoencoder.pt" |
| }, |
| { |
| "_target_": "StatsHandler", |
| "tag_name": "train_loss", |
| "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]" |
| }, |
| { |
| "_target_": "TensorBoardStatsHandler", |
| "log_dir": "@tf_dir", |
| "tag_name": "train_loss", |
| "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]" |
| } |
| ], |
| "trainer": { |
| "_target_": "scripts.ldm_trainer.VaeGanTrainer", |
| "device": "@device", |
| "max_epochs": 1500, |
| "train_data_loader": "@train#dataloader", |
| "g_network": "@gnetwork", |
| "g_optimizer": "@goptimizer", |
| "g_loss_function": "$functools.partial(scripts.losses.generator_loss, disc_net=@dnetwork, loss_perceptual=@loss_perceptual)", |
| "d_network": "@dnetwork", |
| "d_optimizer": "@doptimizer", |
| "d_loss_function": "$functools.partial(scripts.losses.discriminator_loss, disc_net=@dnetwork)", |
| "d_train_steps": 5, |
| "g_update_latents": true, |
| "latent_shape": "@latent_channels", |
| "key_train_metric": "$None", |
| "train_handlers": "@train#handlers" |
| } |
| }, |
| "initialize": [ |
| "$monai.utils.set_determinism(seed=0)" |
| ], |
| "run": [ |
| "$@train#trainer.run()" |
| ] |
| } |
|
|
