From 56a8a56ee44c311bb8b3c59a2235738140502d78 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Wed, 21 Aug 2024 22:51:02 +0000 Subject: [PATCH] Update brats_mri_generative_diffusion to monai1.4 Signed-off-by: Can-Zhao --- .../configs/inference.json | 21 +++++++++++-------- .../configs/inference_autoencoder.json | 9 ++++---- .../configs/metadata.json | 15 ++++++------- .../configs/train_autoencoder.json | 13 ++++++------ .../configs/train_diffusion.json | 12 ++++++----- .../scripts/losses.py | 2 +- 6 files changed, 40 insertions(+), 32 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 6e02b0bd..445803ae 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -20,12 +20,12 @@ 28 ], "autoencoder_def": { - "_target_": "generative.networks.nets.AutoencoderKL", + "_target_": "monai.networks.nets.autoencoderkl.AutoencoderKL", "spatial_dims": "@spatial_dims", "in_channels": "@image_channels", "out_channels": "@image_channels", "latent_channels": "@latent_channels", - "num_channels": [ + "channels": [ 64, 128, 256 @@ -39,14 +39,15 @@ false ], "with_encoder_nonlocal_attn": false, - "with_decoder_nonlocal_attn": false + "with_decoder_nonlocal_attn": false, + "include_fc": false }, "network_def": { - "_target_": "generative.networks.nets.DiffusionModelUNet", + "_target_": "monai.networks.nets.diffusion_model_unet.DiffusionModelUNet", "spatial_dims": "@spatial_dims", "in_channels": "@latent_channels", "out_channels": "@latent_channels", - "num_channels": [ + "channels": [ 256, 256, 512 @@ -61,16 +62,18 @@ 64, 64 ], - "num_res_blocks": 2 + "num_res_blocks": 2, + "include_fc": false, + "use_combined_linear": false }, "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", - "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", + "load_autoencoder": "$@autoencoder_def.load_old_state_dict(torch.load(@load_autoencoder_path))", "autoencoder": "$@autoencoder_def.to(@device)", "load_diffusion_path": "$@model_dir + '/model.pt'", - "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))", + "load_diffusion": "$@network_def.load_old_state_dict(torch.load(@load_diffusion_path))", "diffusion": "$@network_def.to(@device)", "noise_scheduler": { - "_target_": "generative.networks.schedulers.DDIMScheduler", + "_target_": "monai.networks.schedulers.ddim.DDIMScheduler", "_requires_": [ "@load_diffusion", "@load_autoencoder" diff --git a/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json index eb66dee2..27f14ead 100644 --- a/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json @@ -27,12 +27,12 @@ 112 ], "autoencoder_def": { - "_target_": "generative.networks.nets.AutoencoderKL", + "_target_": "monai.networks.nets.autoencoderkl.AutoencoderKL", "spatial_dims": "@spatial_dims", "in_channels": "@image_channels", "out_channels": "@image_channels", "latent_channels": "@latent_channels", - "num_channels": [ + "channels": [ 64, 128, 256 @@ -46,10 +46,11 @@ false ], "with_encoder_nonlocal_attn": false, - "with_decoder_nonlocal_attn": false + "with_decoder_nonlocal_attn": false, + "include_fc": false }, "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", - "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", + "load_autoencoder": "$@autoencoder_def.load_old_state_dict(torch.load(@load_autoencoder_path))", "autoencoder": "$@autoencoder_def.to(@device)", "preprocessing_transforms": [ { diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 2fe64f9c..67cf328e 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -1,7 +1,8 @@ { - "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20230507.json", - "version": "1.0.9", + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json", + "version": "1.0.10", "changelog": { + "1.0.10": "update to use monai 1.4, model ckpt not changed, rm GenerativeAI repo", "1.0.9": "update to use monai 1.3.1", "1.0.8": "update run section", "1.0.7": "update with EnsureChannelFirstd", @@ -13,14 +14,14 @@ "1.0.1": "update dependency, update trained model weights", "1.0.0": "Initial release" }, - "monai_version": "1.3.1", + "monai_version": "1.4.0", "pytorch_version": "2.2.2", "numpy_version": "1.24.4", - "optional_packages_version": { + "required_packages_version": { "nibabel": "5.2.1", - "lpips": "0.1.4", - "monai-generative": "0.2.2" + "lpips": "0.1.4" }, + "supported_apps": {}, "name": "BraTS MRI image latent diffusion generation", "task": "BraTS MRI image synthesis", "description": "A generative model for creating 3D brain MRI from Gaussian noise based on BraTS dataset", @@ -73,7 +74,7 @@ } } }, - "generator_data_format": { + "network_data_format": { "inputs": { "latent": { "type": "noise", diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index ddaa8c41..924f2a78 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -28,21 +28,21 @@ "image_channels": 1, "latent_channels": 8, "discriminator_def": { - "_target_": "generative.networks.nets.PatchDiscriminator", + "_target_": "monai.networks.nets.patchgan_discriminator.PatchDiscriminator", "spatial_dims": "@spatial_dims", "num_layers_d": 3, - "num_channels": 32, + "channels": 32, "in_channels": 1, "out_channels": 1, "norm": "INSTANCE" }, "autoencoder_def": { - "_target_": "generative.networks.nets.AutoencoderKL", + "_target_": "monai.networks.nets.autoencoderkl.AutoencoderKL", "spatial_dims": "@spatial_dims", "in_channels": "@image_channels", "out_channels": "@image_channels", "latent_channels": "@latent_channels", - "num_channels": [ + "channels": [ 64, 128, 256 @@ -56,10 +56,11 @@ false ], "with_encoder_nonlocal_attn": false, - "with_decoder_nonlocal_attn": false + "with_decoder_nonlocal_attn": false, + "include_fc": false }, "perceptual_loss_def": { - "_target_": "generative.losses.PerceptualLoss", + "_target_": "monai.losses.perceptual.PerceptualLoss", "spatial_dims": "@spatial_dims", "network_type": "resnet50", "is_fake_3d": true, diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 85c8ca8a..e58b1618 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -14,14 +14,14 @@ 28 ], "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", - "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", + "load_autoencoder": "$@autoencoder_def.load_old_state_dict(torch.load(@load_autoencoder_path))", "autoencoder": "$@autoencoder_def.to(@device)", "network_def": { "_target_": "generative.networks.nets.DiffusionModelUNet", "spatial_dims": "@spatial_dims", "in_channels": "@latent_channels", "out_channels": "@latent_channels", - "num_channels": [ + "channels": [ 256, 256, 512 @@ -36,7 +36,9 @@ 64, 64 ], - "num_res_blocks": 2 + "num_res_blocks": 2, + "include_fc": false, + "use_combined_linear": false }, "diffusion": "$@network_def.to(@device)", "optimizer": { @@ -55,7 +57,7 @@ }, "scale_factor": "$scripts.utils.compute_scale_factor(@autoencoder,@train#dataloader,@device)", "noise_scheduler": { - "_target_": "generative.networks.schedulers.DDPMScheduler", + "_target_": "monai.networks.schedulers.ddpm.DDPMScheduler", "_requires_": [ "@load_autoencoder" ], @@ -69,7 +71,7 @@ }, "train": { "inferer": { - "_target_": "generative.inferers.LatentDiffusionInferer", + "_target_": "monai.inferers.LatentDiffusionInferer", "scheduler": "@noise_scheduler", "scale_factor": "@scale_factor" }, diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 43536067..1b7c939e 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -9,7 +9,7 @@ # See the License for the specific language governing permissions and import torch -from generative.losses import PatchAdversarialLoss +from monai.losses.adversarial_loss import PatchAdversarialLoss intensity_loss = torch.nn.L1Loss() adv_loss = PatchAdversarialLoss(criterion="least_squares")