From 6bab90c16cbd9183bf57a6747cfbebe0a4b18b24 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 15:18:08 -0700 Subject: [PATCH 001/130] add train config for autoencoder for diffusion bundle Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/LICENSE | 201 ++++++++++++++++++ .../configs/inference.json | 94 ++++++++ .../configs/logging.conf | 21 ++ .../configs/metadata.json | 66 ++++++ .../configs/train_autoencoder.json | 185 ++++++++++++++++ .../docs/README.md | 60 ++++++ .../docs/data_license.txt | 49 +++++ .../scripts/__init__.py | 2 + .../scripts/losses.py | 40 ++++ .../scripts/vae_gan_trainer.py | 201 ++++++++++++++++++ 10 files changed, 919 insertions(+) create mode 100644 models/brats_mri_generative_diffusion/LICENSE create mode 100644 models/brats_mri_generative_diffusion/configs/inference.json create mode 100644 models/brats_mri_generative_diffusion/configs/logging.conf create mode 100644 models/brats_mri_generative_diffusion/configs/metadata.json create mode 100644 models/brats_mri_generative_diffusion/configs/train_autoencoder.json create mode 100644 models/brats_mri_generative_diffusion/docs/README.md create mode 100644 models/brats_mri_generative_diffusion/docs/data_license.txt create mode 100644 models/brats_mri_generative_diffusion/scripts/__init__.py create mode 100644 models/brats_mri_generative_diffusion/scripts/losses.py create mode 100644 models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py diff --git a/models/brats_mri_generative_diffusion/LICENSE b/models/brats_mri_generative_diffusion/LICENSE new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/models/brats_mri_generative_diffusion/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json new file mode 100644 index 00000000..ee3b04fb --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -0,0 +1,94 @@ +{ + "imports": [ + "$import glob" + ], + "bundle_root": ".", + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "ckpt_path": "$@bundle_root + '/models/model.pt'", + "output_dir": "./outputs", + "latent_size": 64, + "num_samples": 10, + "network_def": { + "_target_": "Generator", + "latent_shape": "@latent_size", + "start_shape": [ + 64, + 8, + 8 + ], + "channels": [ + 32, + 16, + 8, + 1 + ], + "strides": [ + 2, + 2, + 2, + 1 + ] + }, + "network": "$@network_def.to(@device)", + "dataset": { + "_target_": "Dataset", + "data": "$[torch.rand(@latent_size) for i in range(@num_samples)]" + }, + "dataloader": { + "_target_": "DataLoader", + "dataset": "@dataset", + "batch_size": 1, + "shuffle": false, + "num_workers": 0 + }, + "inferer": { + "_target_": "SimpleInferer" + }, + "postprocessing": { + "_target_": "Compose", + "transforms": [ + { + "_target_": "Activationsd", + "keys": "pred", + "sigmoid": true + }, + { + "_target_": "ToTensord", + "keys": "pred", + "track_meta": false + }, + { + "_target_": "SaveImaged", + "keys": "pred", + "output_dir": "@output_dir", + "output_ext": "png", + "separate_folder": false, + "scale": 255, + "output_dtype": "$np.uint8", + "meta_key_postfix": null + } + ] + }, + "handlers": [ + { + "_target_": "CheckpointLoader", + "load_path": "@ckpt_path", + "load_dict": { + "model": "@network" + } + } + ], + "evaluator": { + "_target_": "SupervisedEvaluator", + "device": "@device", + "val_data_loader": "@dataloader", + "network": "@network", + "inferer": "@inferer", + "postprocessing": "@postprocessing", + "prepare_batch": "$lambda batchdata, *_,**__: (batchdata.to(@device),None,(),{})", + "val_handlers": "@handlers" + }, + "inferring": [ + "$@evaluator.run()" + ] +} diff --git a/models/brats_mri_generative_diffusion/configs/logging.conf b/models/brats_mri_generative_diffusion/configs/logging.conf new file mode 100644 index 00000000..91c1a21c --- /dev/null +++ b/models/brats_mri_generative_diffusion/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 diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json new file mode 100644 index 00000000..c013d46c --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -0,0 +1,66 @@ +{ + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_20220718.json", + "version": "0.4.2", + "changelog": { + "0.4.2": "add name tag", + "0.4.1": "fix license Copyright error", + "0.4.0": "update license files", + "0.3.0": "Update for 1.0", + "0.2.0": "Unify naming", + "0.1.0": "Initial version" + }, + "monai_version": "1.0.0", + "pytorch_version": "1.10.0", + "numpy_version": "1.21.0", + "optional_packages_version": { + "pytorch-ignite": "0.4.8", + "pillow": "8.4.0" + }, + "name": "MedNIST GAN", + "task": "Generate random hand images from the MedNIST dataset", + "description": "This example of a GAN generator produces hand xray images like those in the MedNIST dataset", + "authors": "MONAI Team", + "copyright": "Copyright (c) MONAI Consortium", + "intended_use": "This is an example of a GAN with generator discriminator networks using MONAI, suitable for demonstration purposes only.", + "data_source": "MedNIST dataset kindly made available by Dr. Bradley J. Erickson M.D., Ph.D. (Department of Radiology, Mayo Clinic)", + "data_type": "jpeg", + "network_data_format": { + "inputs": { + "latent": { + "type": "tuples", + "format": "latent", + "num_channels": 0, + "spatial_shape": [ + 64 + ], + "dtype": "float32", + "value_range": [ + 0, + 1 + ], + "is_patch_data": false, + "channel_def": {} + } + }, + "outputs": { + "pred": { + "type": "image", + "format": "magnitude", + "num_channels": 1, + "spatial_shape": [ + 64, + 64 + ], + "dtype": "float32", + "value_range": [ + 0, + 1 + ], + "is_patch_data": false, + "channel_def": { + "0": "image" + } + } + } + } +} diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json new file mode 100644 index 00000000..fc4433a6 --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -0,0 +1,185 @@ +{ + "imports": [ + "$import functools", + "$import glob", + "$import scripts" + ], + "bundle_root": ".", + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "ckpt_path": "$@bundle_root + '/models/model.pt'", + "output_dir": "$@bundle_root + '/eval'", + "dataset_dir": "/workspace/dataset", + "channel": 0, + "spacing": [1.1, 1.1, 1.1], + "spatial_dims": 3, + "image_channels": 1, + "latent_channels": 8, + "discriminator_def": { + "_target_": "generative.networks.nets.PatchDiscriminator", + "spatial_dims": "@spatial_dims", + "num_layers_d":3, + "num_channels":32, + "in_channels":1, + "out_channels":1, + "norm": "INSTANCE" + }, + "autoencoder_def": { + "_target_": "generative.networks.nets.AutoencoderKL", + "spatial_dims": "@spatial_dims", + "in_channels": "@image_channels", + "out_channels": "@image_channels", + "latent_channels": "@latent_channels", + "num_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 + }, + "perceptual_loss_def": { + "_target_": "generative.losses.PerceptualLoss", + "spatial_dims": "@spatial_dims", + "network_type":"squeeze", + "is_fake_3d":true, + "fake_3d_ratio":0.2 + }, + "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": 1e-5 + }, + "goptimizer": { + "_target_": "torch.optim.Adam", + "params": "$@gnetwork.parameters()", + "lr": 1e-5 + }, + "train": { + "preprocessing_transforms": [ + { + "_target_": "LoadImaged", + "keys": "image" + }, + { + "_target_": "EnsureChannelFirstd", + "keys": "image" + }, + { + "_target_": "Lambdad", + "keys": "image", + "func": "$lambda x: x[@channel, :, :, :]" + }, + { + "_target_": "AddChanneld", + "keys": "image" + }, + { + "_target_": "EnsureTyped", + "keys": "image" + }, + { + "_target_": "Orientationd", + "keys": "image", + "axcodes": "RAS" + }, + { + "_target_": "Spacingd", + "keys": "image", + "pixdim": "@spacing", + "mode": "bilinear" + } + ], + "random_transforms": [ + { + "_target_": "RandSpatialCropd", + "keys": "image", + "roi_size": [112,128,80], + "random_size": false + } + ], + "final_transforms": [ + { + "_target_": "ScaleIntensityRangePercentilesd", + "keys": "image", + "lower":0, + "upper":99.5, + "b_min":0, + "b_max":1 + } + ], + "preprocessing": { + "_target_": "Compose", + "transforms": "$@train#preprocessing_transforms + @train#random_transforms + @train#final_transforms" + }, + "real_dataset": { + "_target_": "monai.apps.DecathlonDataset", + "root_dir": "@dataset_dir", + "task":"Task01_BrainTumour", + "section":"training", + "cache_rate":0, + "num_workers":8, + "download":false, + "transform":"@train#preprocessing" + }, + "real_dataloader": { + "_target_": "DataLoader", + "dataset": "@train#real_dataset", + "batch_size": 2, + "shuffle": true, + "num_workers": 0 + }, + "handlers": [ + { + "_target_": "CheckpointSaver", + "save_dir": "$@bundle_root + '/models'", + "save_dict": { + "model": "@gnetwork" + }, + "save_interval": 0, + "save_final": true, + "epoch_level": true + }, + { + "_target_": "TensorBoardStatsHandler", + "log_dir": "@output_dir", + "tag_name": "train_loss", + "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]" + } + ], + "trainer": { + "_target_": "scripts.vae_gan_trainer.VaeGanTrainer", + "device": "@device", + "max_epochs": 1500, + "train_data_loader": "@train#real_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()" + ] +} diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md new file mode 100644 index 00000000..50b24a76 --- /dev/null +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -0,0 +1,60 @@ + +# Brats MRI 3D Latent Diffusion Generative Model + +This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). + +This is a demonstration network meant to just show the training process for this sort of network with MONAI. + +### Downloading the Dataset +The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. + +Target: 3 tumor subregions +Task: Segmentation +Modality: MRI +Size: 285 3D volumes (4 channels each) +The provided labelled data was partitioned, based on our own split, into training (200 studies), validation (42 studies) and testing (43 studies) datasets. + +### Training + +Assuming the current directory is the bundle directory, and the dataset was extracted to the directory `./MedNIST`, the following command will train the network for 50 epochs: + +``` +python -m monai.bundle run --config_file configs/train.json +``` + + +### Inference + + + +### Export + +The generator can be exported to a Torchscript bundle with the following: + +``` +python -m monai.bundle ckpt_export network_def --filepath mednist_gan.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json +``` + +The model can be loaded without MONAI code after this operation. For example, an image can be generated from a set of random values with: + +```python +import torch +net = torch.jit.load("mednist_gan.ts") +latent = torch.rand(1, 64) +img = net(latent) # (1,1,64,64) +``` + +# License +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. diff --git a/models/brats_mri_generative_diffusion/docs/data_license.txt b/models/brats_mri_generative_diffusion/docs/data_license.txt new file mode 100644 index 00000000..d3d7e227 --- /dev/null +++ b/models/brats_mri_generative_diffusion/docs/data_license.txt @@ -0,0 +1,49 @@ +Third Party Licenses +----------------------------------------------------------------------- + +/*********************************************************************/ +i. Multimodal Brain Tumor Segmentation Challenge 2018 + https://www.med.upenn.edu/sbia/brats2018/data.html +/*********************************************************************/ + +Data Usage Agreement / Citations + +You are free to use and/or refer to the BraTS datasets in your own +research, provided that you always cite the following two manuscripts: + +[1] Menze BH, Jakab A, Bauer S, Kalpathy-Cramer J, Farahani K, Kirby +[J, Burren Y, Porz N, Slotboom J, Wiest R, Lanczi L, Gerstner E, Weber +[MA, Arbel T, Avants BB, Ayache N, Buendia P, Collins DL, Cordier N, +[Corso JJ, Criminisi A, Das T, Delingette H, Demiralp Γ, Durst CR, +[Dojat M, Doyle S, Festa J, Forbes F, Geremia E, Glocker B, Golland P, +[Guo X, Hamamci A, Iftekharuddin KM, Jena R, John NM, Konukoglu E, +[Lashkari D, Mariz JA, Meier R, Pereira S, Precup D, Price SJ, Raviv +[TR, Reza SM, Ryan M, Sarikaya D, Schwartz L, Shin HC, Shotton J, +[Silva CA, Sousa N, Subbanna NK, Szekely G, Taylor TJ, Thomas OM, +[Tustison NJ, Unal G, Vasseur F, Wintermark M, Ye DH, Zhao L, Zhao B, +[Zikic D, Prastawa M, Reyes M, Van Leemput K. "The Multimodal Brain +[Tumor Image Segmentation Benchmark (BRATS)", IEEE Transactions on +[Medical Imaging 34(10), 1993-2024 (2015) DOI: +[10.1109/TMI.2014.2377694 + +[2] Bakas S, Akbari H, Sotiras A, Bilello M, Rozycki M, Kirby JS, +[Freymann JB, Farahani K, Davatzikos C. "Advancing The Cancer Genome +[Atlas glioma MRI collections with expert segmentation labels and +[radiomic features", Nature Scientific Data, 4:170117 (2017) DOI: +[10.1038/sdata.2017.117 + +In addition, if there are no restrictions imposed from the +journal/conference you submit your paper about citing "Data +Citations", please be specific and also cite the following: + +[3] Bakas S, Akbari H, Sotiras A, Bilello M, Rozycki M, Kirby J, +[Freymann J, Farahani K, Davatzikos C. "Segmentation Labels and +[Radiomic Features for the Pre-operative Scans of the TCGA-GBM +[collection", The Cancer Imaging Archive, 2017. DOI: +[10.7937/K9/TCIA.2017.KLXWJJ1Q + +[4] Bakas S, Akbari H, Sotiras A, Bilello M, Rozycki M, Kirby J, +[Freymann J, Farahani K, Davatzikos C. "Segmentation Labels and +[Radiomic Features for the Pre-operative Scans of the TCGA-LGG +[collection", The Cancer Imaging Archive, 2017. DOI: +[10.7937/K9/TCIA.2017.GJQ7R0EF diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py new file mode 100644 index 00000000..dd4c73f6 --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -0,0 +1,2 @@ +from . import losses +from . import vae_gan_trainer \ No newline at end of file diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py new file mode 100644 index 00000000..26cf37e2 --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -0,0 +1,40 @@ +import torch +from generative.losses import PatchAdversarialLoss, PerceptualLoss +from generative.networks.nets import PatchDiscriminator + +intensity_loss = torch.nn.L1Loss() +adv_loss = PatchAdversarialLoss(criterion="least_squares") +loss_perceptual = PerceptualLoss(spatial_dims=3, network_type="squeeze", is_fake_3d=True, fake_3d_ratio=0.2) + +adv_weight = 0.01 +perceptual_weight = 0.01 +# kl_weight: important hyper-parameter. +# If too large, decoder cannot recon good results from latent space. +# If too small, latent space will not be regularized enough for the diffusion model +kl_weight = 1e-7 + +def KL_loss(z_mu, z_sigma): + kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) + return torch.sum(kl_loss) / kl_loss.shape[0] + +def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): + recons_loss = intensity_loss(gen_images, real_images) + kl_loss = KL_loss(z_mu, z_sigma) + p_loss = loss_perceptual(gen_images.float(), real_images.float()) + loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss + + logits_fake = disc_net(gen_images)[-1] + generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False) + loss_g = loss_g + adv_weight * generator_loss + + return loss_g + + +def discriminator_loss(gen_images, real_images, disc_net): + logits_fake = disc_net(gen_images.contiguous().detach())[-1] + loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True) + logits_real = disc_net(real_images.contiguous().detach())[-1] + loss_d_real = adv_loss(logits_real, target_is_real=True, for_discriminator=True) + discriminator_loss = (loss_d_fake + loss_d_real) * 0.5 + loss_d = adv_weight * discriminator_loss + return loss_d diff --git a/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py b/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py new file mode 100644 index 00000000..d7c2da01 --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py @@ -0,0 +1,201 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Sequence + +import torch +from torch.optim.optimizer import Optimizer +from torch.utils.data import DataLoader + +from monai.config import IgniteInfo +from monai.engines.utils import default_make_latent, default_metric_cmp_fn, default_prepare_batch +from monai.inferers import Inferer, SimpleInferer +from monai.transforms import Transform +from monai.utils import GanKeys, min_version, optional_import +from monai.utils.enums import CommonKeys, GanKeys + +if TYPE_CHECKING: + from ignite.engine import Engine, EventEnum + from ignite.metrics import Metric +else: + Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") + Metric, _ = optional_import("ignite.metrics", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Metric") + EventEnum, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "EventEnum") +from monai.engines.trainer import Trainer + +class VaeGanTrainer(Trainer): + """ + Generative adversarial network training based on Goodfellow et al. 2014 https://arxiv.org/abs/1406.266, + inherits from ``Trainer`` and ``Workflow``. + Training Loop: for each batch of data size `m` + 1. Generate `m` fakes from random latent codes. + 2. Update discriminator with these fakes and current batch reals, repeated d_train_steps times. + 3. If g_update_latents, generate `m` fakes from new random latent codes. + 4. Update generator with these fakes using discriminator feedback. + Args: + device: an object representing the device on which to run. + max_epochs: the total epoch number for engine to run. + train_data_loader: Core ignite engines uses `DataLoader` for training loop batchdata. + g_network: generator (G) network architecture. + g_optimizer: G optimizer function. + g_loss_function: G loss function for optimizer. + d_network: discriminator (D) network architecture. + d_optimizer: D optimizer function. + d_loss_function: D loss function for optimizer. + epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`. + g_inferer: inference method to execute G model forward. Defaults to ``SimpleInferer()``. + d_inferer: inference method to execute D model forward. Defaults to ``SimpleInferer()``. + d_train_steps: number of times to update D with real data minibatch. Defaults to ``1``. + latent_shape: size of G input latent code. Defaults to ``64``. + non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously + with respect to the host. For other cases, this argument has no effect. + d_prepare_batch: callback function to prepare batchdata for D inferer. + Defaults to return ``GanKeys.REALS`` in batchdata dict. for more details please refer to: + https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. + g_prepare_batch: callback function to create batch of latent input for G inferer. + Defaults to return random latents. for more details please refer to: + https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. + g_update_latents: Calculate G loss with new latent codes. Defaults to ``True``. + iteration_update: the callable function for every iteration, expect to accept `engine` + and `engine.state.batch` as inputs, return data will be stored in `engine.state.output`. + if not provided, use `self._iteration()` instead. for more details please refer to: + https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html. + postprocessing: execute additional transformation for the model output data. + Typically, several Tensor based transforms composed by `Compose`. + key_train_metric: compute metric when every iteration completed, and save average value to + engine.state.metrics when epoch completed. key_train_metric is the main metric to compare and save the + checkpoint into files. + additional_metrics: more Ignite metrics that also attach to Ignite Engine. + metric_cmp_fn: function to compare current key metric with previous best key metric value, + it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update + `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. + train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: + CheckpointHandler, StatsHandler, etc. + decollate: whether to decollate the batch-first data to a list of data after model computation, + recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. + default to `True`. + optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. + more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. + """ + + def __init__( + self, + device: str | torch.device, + max_epochs: int, + train_data_loader: DataLoader, + g_network: torch.nn.Module, + g_optimizer: Optimizer, + g_loss_function: Callable, + d_network: torch.nn.Module, + d_optimizer: Optimizer, + d_loss_function: Callable, + epoch_length: int | None = None, + g_inferer: Inferer | None = None, + d_inferer: Inferer | None = None, + d_train_steps: int = 1, + latent_shape: int = 64, + non_blocking: bool = False, + d_prepare_batch: Callable = default_prepare_batch, + g_prepare_batch: Callable = default_prepare_batch, + g_update_latents: bool = True, + iteration_update: Callable[[Engine, Any], Any] | None = None, + postprocessing: Transform | None = None, + key_train_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, + metric_cmp_fn: Callable = default_metric_cmp_fn, + train_handlers: Sequence | None = None, + decollate: bool = True, + optim_set_to_none: bool = False, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, + ): + if not isinstance(train_data_loader, DataLoader): + raise ValueError("train_data_loader must be PyTorch DataLoader.") + + # set up Ignite engine and environments + super().__init__( + device=device, + max_epochs=max_epochs, + data_loader=train_data_loader, + epoch_length=epoch_length, + non_blocking=non_blocking, + prepare_batch=d_prepare_batch, + iteration_update=iteration_update, + key_metric=key_train_metric, + additional_metrics=additional_metrics, + metric_cmp_fn=metric_cmp_fn, + handlers=train_handlers, + postprocessing=postprocessing, + decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, + ) + self.g_network = g_network + self.g_optimizer = g_optimizer + self.g_loss_function = g_loss_function + self.g_inferer = SimpleInferer() if g_inferer is None else g_inferer + self.d_network = d_network + self.d_optimizer = d_optimizer + self.d_loss_function = d_loss_function + self.d_inferer = SimpleInferer() if d_inferer is None else d_inferer + self.d_train_steps = d_train_steps + self.latent_shape = latent_shape + self.g_prepare_batch = g_prepare_batch + self.g_update_latents = g_update_latents + self.optim_set_to_none = optim_set_to_none + + def _iteration( + self, engine: VaeGanTrainer, batchdata: dict | Sequence + ) -> dict[str, torch.Tensor | int | float | bool]: + """ + Callback function for Adversarial Training processing logic of 1 iteration in Ignite Engine. + Args: + engine: `VaeGanTrainer` to execute operation for an iteration. + batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. + Raises: + ValueError: must provide batch data for current iteration. + """ + if batchdata is None: + raise ValueError("must provide batch data for current iteration.") + + d_input = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)[0] + batch_size = engine.data_loader.batch_size # type: ignore + g_input = d_input + g_output, z_mu, z_sigma = engine.g_inferer(g_input, engine.g_network) + + # Train Discriminator + d_total_loss = torch.zeros(1) + for _ in range(engine.d_train_steps): + engine.d_optimizer.zero_grad(set_to_none=engine.optim_set_to_none) + dloss = engine.d_loss_function(g_output, d_input) + dloss.backward() + engine.d_optimizer.step() + d_total_loss += dloss.item() + + # Train Generator + engine.g_optimizer.zero_grad(set_to_none=engine.optim_set_to_none) + g_loss = engine.g_loss_function(g_output, g_input, z_mu, z_sigma) + g_loss.backward() + engine.g_optimizer.step() + + return { + GanKeys.REALS: d_input, + GanKeys.FAKES: g_output, + GanKeys.LATENTS: g_input, + GanKeys.GLOSS: g_loss.item(), + GanKeys.DLOSS: d_total_loss.item(), + } \ No newline at end of file From 4bf387f3337cafa3e6635abc7d8f8cd28cba9861 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Apr 2023 22:18:51 +0000 Subject: [PATCH 002/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../configs/train_autoencoder.json | 53 +++++++++++-------- .../docs/README.md | 2 +- .../scripts/__init__.py | 2 +- .../scripts/losses.py | 8 +-- .../scripts/vae_gan_trainer.py | 4 +- 5 files changed, 38 insertions(+), 31 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index fc4433a6..c825e62f 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -10,17 +10,21 @@ "output_dir": "$@bundle_root + '/eval'", "dataset_dir": "/workspace/dataset", "channel": 0, - "spacing": [1.1, 1.1, 1.1], + "spacing": [ + 1.1, + 1.1, + 1.1 + ], "spatial_dims": 3, "image_channels": 1, "latent_channels": 8, "discriminator_def": { "_target_": "generative.networks.nets.PatchDiscriminator", "spatial_dims": "@spatial_dims", - "num_layers_d":3, - "num_channels":32, - "in_channels":1, - "out_channels":1, + "num_layers_d": 3, + "num_channels": 32, + "in_channels": 1, + "out_channels": 1, "norm": "INSTANCE" }, "autoencoder_def": { @@ -48,9 +52,9 @@ "perceptual_loss_def": { "_target_": "generative.losses.PerceptualLoss", "spatial_dims": "@spatial_dims", - "network_type":"squeeze", - "is_fake_3d":true, - "fake_3d_ratio":0.2 + "network_type": "squeeze", + "is_fake_3d": true, + "fake_3d_ratio": 0.2 }, "dnetwork": "$@discriminator_def.to(@device)", "gnetwork": "$@autoencoder_def.to(@device)", @@ -58,13 +62,13 @@ "doptimizer": { "_target_": "torch.optim.Adam", "params": "$@dnetwork.parameters()", - "lr": 1e-5 + "lr": 1e-05 }, "goptimizer": { "_target_": "torch.optim.Adam", "params": "$@gnetwork.parameters()", - "lr": 1e-5 - }, + "lr": 1e-05 + }, "train": { "preprocessing_transforms": [ { @@ -104,7 +108,11 @@ { "_target_": "RandSpatialCropd", "keys": "image", - "roi_size": [112,128,80], + "roi_size": [ + 112, + 128, + 80 + ], "random_size": false } ], @@ -112,10 +120,10 @@ { "_target_": "ScaleIntensityRangePercentilesd", "keys": "image", - "lower":0, - "upper":99.5, - "b_min":0, - "b_max":1 + "lower": 0, + "upper": 99.5, + "b_min": 0, + "b_max": 1 } ], "preprocessing": { @@ -125,12 +133,12 @@ "real_dataset": { "_target_": "monai.apps.DecathlonDataset", "root_dir": "@dataset_dir", - "task":"Task01_BrainTumour", - "section":"training", - "cache_rate":0, - "num_workers":8, - "download":false, - "transform":"@train#preprocessing" + "task": "Task01_BrainTumour", + "section": "training", + "cache_rate": 0, + "num_workers": 8, + "download": false, + "transform": "@train#preprocessing" }, "real_dataloader": { "_target_": "DataLoader", @@ -175,7 +183,6 @@ "train_handlers": "@train#handlers" } }, - "initialize": [ "$monai.utils.set_determinism(seed=0)" ], diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 50b24a76..26ac9189 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index dd4c73f6..dd12b6cf 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -1,2 +1,2 @@ from . import losses -from . import vae_gan_trainer \ No newline at end of file +from . import vae_gan_trainer diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 26cf37e2..4820d82f 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -26,11 +26,11 @@ def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_percep logits_fake = disc_net(gen_images)[-1] generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False) loss_g = loss_g + adv_weight * generator_loss - + return loss_g - - -def discriminator_loss(gen_images, real_images, disc_net): + + +def discriminator_loss(gen_images, real_images, disc_net): logits_fake = disc_net(gen_images.contiguous().detach())[-1] loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True) logits_real = disc_net(real_images.contiguous().detach())[-1] diff --git a/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py b/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py index d7c2da01..d5a6b6eb 100644 --- a/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py +++ b/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py @@ -171,7 +171,7 @@ def _iteration( """ if batchdata is None: raise ValueError("must provide batch data for current iteration.") - + d_input = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)[0] batch_size = engine.data_loader.batch_size # type: ignore g_input = d_input @@ -198,4 +198,4 @@ def _iteration( GanKeys.LATENTS: g_input, GanKeys.GLOSS: g_loss.item(), GanKeys.DLOSS: d_total_loss.item(), - } \ No newline at end of file + } From 4b3335b86a5fbbfe07cc1425f0c55544ef360d9c Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 18:21:52 -0700 Subject: [PATCH 003/130] add train config for diffusion model for diffusion bundle Signed-off-by: Can-Zhao --- .../configs/inference.json | 94 ---------- .../configs/metadata.json | 62 ++++--- .../configs/train_autoencoder.json | 161 +++++++++--------- .../configs/train_diffusion.json | 120 +++++++++++++ .../docs/README.md | 8 +- 5 files changed, 233 insertions(+), 212 deletions(-) delete mode 100644 models/brats_mri_generative_diffusion/configs/inference.json create mode 100644 models/brats_mri_generative_diffusion/configs/train_diffusion.json diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json deleted file mode 100644 index ee3b04fb..00000000 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ /dev/null @@ -1,94 +0,0 @@ -{ - "imports": [ - "$import glob" - ], - "bundle_root": ".", - "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", - "ckpt_path": "$@bundle_root + '/models/model.pt'", - "output_dir": "./outputs", - "latent_size": 64, - "num_samples": 10, - "network_def": { - "_target_": "Generator", - "latent_shape": "@latent_size", - "start_shape": [ - 64, - 8, - 8 - ], - "channels": [ - 32, - 16, - 8, - 1 - ], - "strides": [ - 2, - 2, - 2, - 1 - ] - }, - "network": "$@network_def.to(@device)", - "dataset": { - "_target_": "Dataset", - "data": "$[torch.rand(@latent_size) for i in range(@num_samples)]" - }, - "dataloader": { - "_target_": "DataLoader", - "dataset": "@dataset", - "batch_size": 1, - "shuffle": false, - "num_workers": 0 - }, - "inferer": { - "_target_": "SimpleInferer" - }, - "postprocessing": { - "_target_": "Compose", - "transforms": [ - { - "_target_": "Activationsd", - "keys": "pred", - "sigmoid": true - }, - { - "_target_": "ToTensord", - "keys": "pred", - "track_meta": false - }, - { - "_target_": "SaveImaged", - "keys": "pred", - "output_dir": "@output_dir", - "output_ext": "png", - "separate_folder": false, - "scale": 255, - "output_dtype": "$np.uint8", - "meta_key_postfix": null - } - ] - }, - "handlers": [ - { - "_target_": "CheckpointLoader", - "load_path": "@ckpt_path", - "load_dict": { - "model": "@network" - } - } - ], - "evaluator": { - "_target_": "SupervisedEvaluator", - "device": "@device", - "val_data_loader": "@dataloader", - "network": "@network", - "inferer": "@inferer", - "postprocessing": "@postprocessing", - "prepare_batch": "$lambda batchdata, *_,**__: (batchdata.to(@device),None,(),{})", - "val_handlers": "@handlers" - }, - "inferring": [ - "$@evaluator.run()" - ] -} diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index c013d46c..015dc95d 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -1,55 +1,51 @@ { - "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_20220718.json", - "version": "0.4.2", + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json", + "version": "1.0.0", "changelog": { - "0.4.2": "add name tag", - "0.4.1": "fix license Copyright error", - "0.4.0": "update license files", - "0.3.0": "Update for 1.0", - "0.2.0": "Unify naming", - "0.1.0": "Initial version" + "1.0.0": "Initial release" }, - "monai_version": "1.0.0", - "pytorch_version": "1.10.0", - "numpy_version": "1.21.0", + "monai_version": "1.1.0", + "pytorch_version": "1.13.0", + "numpy_version": "1.22.4", "optional_packages_version": { - "pytorch-ignite": "0.4.8", - "pillow": "8.4.0" + "nibabel": "4.0.1", + "generative": "0.1.0" }, - "name": "MedNIST GAN", - "task": "Generate random hand images from the MedNIST dataset", - "description": "This example of a GAN generator produces hand xray images like those in the MedNIST dataset", - "authors": "MONAI Team", + "task": "BraTS MRI image synthesis", + "description": "A generative model for creating 3D brain MRI from Gaussian noise based on BraTS dataset", + "authors": "MONAI team", "copyright": "Copyright (c) MONAI Consortium", - "intended_use": "This is an example of a GAN with generator discriminator networks using MONAI, suitable for demonstration purposes only.", - "data_source": "MedNIST dataset kindly made available by Dr. Bradley J. Erickson M.D., Ph.D. (Department of Radiology, Mayo Clinic)", - "data_type": "jpeg", + "data_source": "https://www.med.upenn.edu/sbia/brats2018/data.html", + "data_type": "nibabel", + "image_classes": "Flair brain MRI with 1.1x1.1x1.1 mm voxel size", + "eval_metrics": { + }, + "intended_use": "This is a research tool/prototype and not to be used clinically", + "references": [], "network_data_format": { "inputs": { - "latent": { - "type": "tuples", - "format": "latent", - "num_channels": 0, - "spatial_shape": [ - 64 - ], + "image": { + "type": "image", + "format": "random noise", + "num_channels": 1, "dtype": "float32", "value_range": [ 0, 1 ], - "is_patch_data": false, - "channel_def": {} + "is_patch_data": false } }, "outputs": { "pred": { "type": "image", "format": "magnitude", + "modality": "MR", "num_channels": 1, "spatial_shape": [ - 64, - 64 + 144, + 176, + 112 ], "dtype": "float32", "value_range": [ @@ -58,9 +54,9 @@ ], "is_patch_data": false, "channel_def": { - "0": "image" + "0": "Flair" } } } } -} +} \ No newline at end of file diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index c825e62f..4cb38d25 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -6,25 +6,23 @@ ], "bundle_root": ".", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", - "ckpt_path": "$@bundle_root + '/models/model.pt'", - "output_dir": "$@bundle_root + '/eval'", + "ckpt_path": "$@bundle_root + '/models/model_autoencoder.pt'", + "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "/workspace/dataset", + "train_batch_size": 2, + "train_patch_size": [112,128,80], "channel": 0, - "spacing": [ - 1.1, - 1.1, - 1.1 - ], + "spacing": [1.1, 1.1, 1.1], "spatial_dims": 3, "image_channels": 1, "latent_channels": 8, "discriminator_def": { "_target_": "generative.networks.nets.PatchDiscriminator", "spatial_dims": "@spatial_dims", - "num_layers_d": 3, - "num_channels": 32, - "in_channels": 1, - "out_channels": 1, + "num_layers_d":3, + "num_channels":32, + "in_channels":1, + "out_channels":1, "norm": "INSTANCE" }, "autoencoder_def": { @@ -52,9 +50,9 @@ "perceptual_loss_def": { "_target_": "generative.losses.PerceptualLoss", "spatial_dims": "@spatial_dims", - "network_type": "squeeze", - "is_fake_3d": true, - "fake_3d_ratio": 0.2 + "network_type":"squeeze", + "is_fake_3d":true, + "fake_3d_ratio":0.2 }, "dnetwork": "$@discriminator_def.to(@device)", "gnetwork": "$@autoencoder_def.to(@device)", @@ -62,88 +60,84 @@ "doptimizer": { "_target_": "torch.optim.Adam", "params": "$@dnetwork.parameters()", - "lr": 1e-05 + "lr": 1e-5 }, "goptimizer": { "_target_": "torch.optim.Adam", "params": "$@gnetwork.parameters()", - "lr": 1e-05 - }, - "train": { - "preprocessing_transforms": [ - { - "_target_": "LoadImaged", - "keys": "image" - }, - { - "_target_": "EnsureChannelFirstd", - "keys": "image" - }, - { - "_target_": "Lambdad", - "keys": "image", - "func": "$lambda x: x[@channel, :, :, :]" - }, - { - "_target_": "AddChanneld", - "keys": "image" - }, - { - "_target_": "EnsureTyped", - "keys": "image" - }, - { - "_target_": "Orientationd", - "keys": "image", - "axcodes": "RAS" - }, - { - "_target_": "Spacingd", - "keys": "image", - "pixdim": "@spacing", - "mode": "bilinear" - } - ], - "random_transforms": [ + "lr": 1e-5 + }, + "preprocessing_transforms": [ + { + "_target_": "LoadImaged", + "keys": "image" + }, + { + "_target_": "EnsureChannelFirstd", + "keys": "image" + }, + { + "_target_": "Lambdad", + "keys": "image", + "func": "$lambda x: x[@channel, :, :, :]" + }, + { + "_target_": "AddChanneld", + "keys": "image" + }, + { + "_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": [ - 112, - 128, - 80 - ], + "roi_size": "@train_patch_size", "random_size": false } - ], - "final_transforms": [ - { - "_target_": "ScaleIntensityRangePercentilesd", - "keys": "image", - "lower": 0, - "upper": 99.5, - "b_min": 0, - "b_max": 1 - } - ], + ], "preprocessing": { "_target_": "Compose", - "transforms": "$@train#preprocessing_transforms + @train#random_transforms + @train#final_transforms" + "transforms": "$@preprocessing_transforms + @train#crop_transforms + @final_transforms" }, - "real_dataset": { + "dataset": { "_target_": "monai.apps.DecathlonDataset", "root_dir": "@dataset_dir", - "task": "Task01_BrainTumour", - "section": "training", - "cache_rate": 0, - "num_workers": 8, - "download": false, - "transform": "@train#preprocessing" + "task":"Task01_BrainTumour", + "section":"training", + "cache_rate":0, + "num_workers":8, + "download":false, + "transform":"@train#preprocessing" }, - "real_dataloader": { + "dataloader": { "_target_": "DataLoader", - "dataset": "@train#real_dataset", - "batch_size": 2, + "dataset": "@train#dataset", + "batch_size": "@train_batch_size", "shuffle": true, "num_workers": 0 }, @@ -160,16 +154,16 @@ }, { "_target_": "TensorBoardStatsHandler", - "log_dir": "@output_dir", + "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.vae_gan_trainer.VaeGanTrainer", + "_target_": "scripts.ldm_trainer.VaeGanTrainer", "device": "@device", "max_epochs": 1500, - "train_data_loader": "@train#real_dataloader", + "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)", @@ -183,6 +177,7 @@ "train_handlers": "@train#handlers" } }, + "initialize": [ "$monai.utils.set_determinism(seed=0)" ], diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json new file mode 100644 index 00000000..3d264f25 --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -0,0 +1,120 @@ +{ + "ckpt_path": "$@bundle_root + '/models/model_ldm.pt'", + "train_batch_size": 4, + "train_patch_size": [144,176,112], + "latent_shape": [8,36,44,28], + "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", + "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", + "autoencoder": "$@autoencoder_def.to(@device)", + "diffusion_def": { + "_target_": "generative.networks.nets.DiffusionModelUNet", + "spatial_dims": "@spatial_dims", + "in_channels": "@latent_channels", + "out_channels": "@latent_channels", + "num_channels":[256, 256, 512], + "attention_levels":[false, true, true], + "num_head_channels":[0, 64, 64], + "num_res_blocks": 2 + }, + "diffusion": "$@diffusion_def.to(@device)", + "optimizer": { + "_target_": "torch.optim.Adam", + "params": "$@diffusion.parameters()", + "lr": 1e-5 + }, + "lr_scheduler": { + "_target_": "torch.optim.lr_scheduler.MultiStepLR", + "optimizer": "@optimizer", + "milestones": [100, 1000], + "gamma": 0.1 + }, + "scale_factor": 1.0, + "noise_scheduler":{ + "_target_": "generative.networks.schedulers.DDPMScheduler", + "beta_schedule": "scaled_linear", + "num_train_timesteps": 1000, + "beta_start": 0.0015, + "beta_end": 0.0195 + }, + "inferer": { + "_target_": "generative.inferers.LatentDiffusionInferer", + "scheduler": "@noise_scheduler", + "scale_factor": "@scale_factor" + }, + "loss": { + "_target_": "torch.nn.MSELoss" + }, + "train": { + "crop_transforms": [ + { + "_target_": "CenterSpatialCropd", + "keys": "image", + "roi_size": "@train_patch_size" + } + ], + "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":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_": "LrScheduleHandler", + "lr_scheduler": "@lr_scheduler", + "print_lr": true + }, + { + "_target_": "CheckpointSaver", + "save_dir": "$@bundle_root + '/models'", + "save_dict": { + "model": "@diffusion" + }, + "save_interval": 0, + "save_final": true, + "epoch_level": true + }, + { + "_target_": "TensorBoardStatsHandler", + "log_dir": "@tf_dir", + "tag_name": "train_diffusion_loss", + "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)" + } + ], + "trainer": { + "_target_": "scripts.ldm_trainer.LDMTrainer", + "device": "@device", + "max_epochs": 10000, + "train_data_loader": "@train#dataloader", + "network": "@diffusion", + "autoencoder_model": "@autoencoder", + "optimizer": "@optimizer", + "loss_function": "@loss", + "latent_shape": "@latent_shape", + "inferer": "@inferer", + "key_train_metric": "$None", + "train_handlers": "@train#handlers" + } + }, + "initialize": [ + "$monai.utils.set_determinism(seed=0)" + ], + "run": [ + "$@train#trainer.run()" + ] +} diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 26ac9189..fe9e7b3b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. @@ -19,7 +19,11 @@ The provided labelled data was partitioned, based on our own split, into trainin Assuming the current directory is the bundle directory, and the dataset was extracted to the directory `./MedNIST`, the following command will train the network for 50 epochs: ``` -python -m monai.bundle run --config_file configs/train.json +python -m monai.bundle run --config_file configs/train_autoencoder.json +``` + +``` +python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` From 79dc1642e6cac620213097aa2bee12a7669bc0b0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 01:22:07 +0000 Subject: [PATCH 004/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../configs/metadata.json | 5 +- .../configs/train_autoencoder.json | 57 +++++++++-------- .../configs/train_diffusion.json | 62 +++++++++++++------ .../docs/README.md | 2 +- 4 files changed, 78 insertions(+), 48 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 015dc95d..3a7f2ccf 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -18,8 +18,7 @@ "data_source": "https://www.med.upenn.edu/sbia/brats2018/data.html", "data_type": "nibabel", "image_classes": "Flair brain MRI with 1.1x1.1x1.1 mm voxel size", - "eval_metrics": { - }, + "eval_metrics": {}, "intended_use": "This is a research tool/prototype and not to be used clinically", "references": [], "network_data_format": { @@ -59,4 +58,4 @@ } } } -} \ No newline at end of file +} diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 4cb38d25..b6fa8d85 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -10,19 +10,27 @@ "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "/workspace/dataset", "train_batch_size": 2, - "train_patch_size": [112,128,80], + "train_patch_size": [ + 112, + 128, + 80 + ], "channel": 0, - "spacing": [1.1, 1.1, 1.1], + "spacing": [ + 1.1, + 1.1, + 1.1 + ], "spatial_dims": 3, "image_channels": 1, "latent_channels": 8, "discriminator_def": { "_target_": "generative.networks.nets.PatchDiscriminator", "spatial_dims": "@spatial_dims", - "num_layers_d":3, - "num_channels":32, - "in_channels":1, - "out_channels":1, + "num_layers_d": 3, + "num_channels": 32, + "in_channels": 1, + "out_channels": 1, "norm": "INSTANCE" }, "autoencoder_def": { @@ -50,9 +58,9 @@ "perceptual_loss_def": { "_target_": "generative.losses.PerceptualLoss", "spatial_dims": "@spatial_dims", - "network_type":"squeeze", - "is_fake_3d":true, - "fake_3d_ratio":0.2 + "network_type": "squeeze", + "is_fake_3d": true, + "fake_3d_ratio": 0.2 }, "dnetwork": "$@discriminator_def.to(@device)", "gnetwork": "$@autoencoder_def.to(@device)", @@ -60,13 +68,13 @@ "doptimizer": { "_target_": "torch.optim.Adam", "params": "$@dnetwork.parameters()", - "lr": 1e-5 + "lr": 1e-05 }, "goptimizer": { "_target_": "torch.optim.Adam", "params": "$@gnetwork.parameters()", - "lr": 1e-5 - }, + "lr": 1e-05 + }, "preprocessing_transforms": [ { "_target_": "LoadImaged", @@ -105,13 +113,13 @@ { "_target_": "ScaleIntensityRangePercentilesd", "keys": "image", - "lower":0, - "upper":99.5, - "b_min":0, - "b_max":1 + "lower": 0, + "upper": 99.5, + "b_min": 0, + "b_max": 1 } ], - "train": { + "train": { "crop_transforms": [ { "_target_": "RandSpatialCropd", @@ -119,7 +127,7 @@ "roi_size": "@train_patch_size", "random_size": false } - ], + ], "preprocessing": { "_target_": "Compose", "transforms": "$@preprocessing_transforms + @train#crop_transforms + @final_transforms" @@ -127,12 +135,12 @@ "dataset": { "_target_": "monai.apps.DecathlonDataset", "root_dir": "@dataset_dir", - "task":"Task01_BrainTumour", - "section":"training", - "cache_rate":0, - "num_workers":8, - "download":false, - "transform":"@train#preprocessing" + "task": "Task01_BrainTumour", + "section": "training", + "cache_rate": 0, + "num_workers": 8, + "download": false, + "transform": "@train#preprocessing" }, "dataloader": { "_target_": "DataLoader", @@ -177,7 +185,6 @@ "train_handlers": "@train#handlers" } }, - "initialize": [ "$monai.utils.set_determinism(seed=0)" ], diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 3d264f25..8e0c2fa5 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -1,8 +1,17 @@ { - "ckpt_path": "$@bundle_root + '/models/model_ldm.pt'", + "ckpt_path": "$@bundle_root + '/models/model_ldm.pt'", "train_batch_size": 4, - "train_patch_size": [144,176,112], - "latent_shape": [8,36,44,28], + "train_patch_size": [ + 144, + 176, + 112 + ], + "latent_shape": [ + 8, + 36, + 44, + 28 + ], "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", "autoencoder": "$@autoencoder_def.to(@device)", @@ -11,30 +20,45 @@ "spatial_dims": "@spatial_dims", "in_channels": "@latent_channels", "out_channels": "@latent_channels", - "num_channels":[256, 256, 512], - "attention_levels":[false, true, true], - "num_head_channels":[0, 64, 64], + "num_channels": [ + 256, + 256, + 512 + ], + "attention_levels": [ + false, + true, + true + ], + "num_head_channels": [ + 0, + 64, + 64 + ], "num_res_blocks": 2 }, "diffusion": "$@diffusion_def.to(@device)", "optimizer": { "_target_": "torch.optim.Adam", "params": "$@diffusion.parameters()", - "lr": 1e-5 + "lr": 1e-05 }, "lr_scheduler": { "_target_": "torch.optim.lr_scheduler.MultiStepLR", "optimizer": "@optimizer", - "milestones": [100, 1000], + "milestones": [ + 100, + 1000 + ], "gamma": 0.1 }, "scale_factor": 1.0, - "noise_scheduler":{ + "noise_scheduler": { "_target_": "generative.networks.schedulers.DDPMScheduler", "beta_schedule": "scaled_linear", "num_train_timesteps": 1000, "beta_start": 0.0015, - "beta_end": 0.0195 + "beta_end": 0.0195 }, "inferer": { "_target_": "generative.inferers.LatentDiffusionInferer", @@ -44,14 +68,14 @@ "loss": { "_target_": "torch.nn.MSELoss" }, - "train": { + "train": { "crop_transforms": [ { "_target_": "CenterSpatialCropd", "keys": "image", "roi_size": "@train_patch_size" } - ], + ], "preprocessing": { "_target_": "Compose", "transforms": "$@preprocessing_transforms + @train#crop_transforms + @final_transforms" @@ -59,12 +83,12 @@ "dataset": { "_target_": "monai.apps.DecathlonDataset", "root_dir": "@dataset_dir", - "task":"Task01_BrainTumour", - "section":"training", - "cache_rate":0, - "num_workers":8, - "download":false, - "transform":"@train#preprocessing" + "task": "Task01_BrainTumour", + "section": "training", + "cache_rate": 0, + "num_workers": 8, + "download": false, + "transform": "@train#preprocessing" }, "dataloader": { "_target_": "DataLoader", @@ -110,7 +134,7 @@ "key_train_metric": "$None", "train_handlers": "@train#handlers" } - }, + }, "initialize": [ "$monai.utils.set_determinism(seed=0)" ], diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index fe9e7b3b..8d4fbda8 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. From d4a81fb52634460ba76bb3614516e069df61e988 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 18:25:28 -0700 Subject: [PATCH 005/130] reformat Signed-off-by: Can-Zhao --- .../scripts/__init__.py | 11 + .../{vae_gan_trainer.py => ldm_trainer.py} | 190 +++++++++++++++++- .../scripts/losses.py | 9 +- 3 files changed, 200 insertions(+), 10 deletions(-) rename models/brats_mri_generative_diffusion/scripts/{vae_gan_trainer.py => ldm_trainer.py} (51%) diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index dd12b6cf..85ee5174 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -1,2 +1,13 @@ +# 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. + from . import losses from . import vae_gan_trainer diff --git a/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py b/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py similarity index 51% rename from models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py rename to models/brats_mri_generative_diffusion/scripts/ldm_trainer.py index d5a6b6eb..16367671 100644 --- a/models/brats_mri_generative_diffusion/scripts/vae_gan_trainer.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py @@ -14,15 +14,14 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence import torch -from torch.optim.optimizer import Optimizer -from torch.utils.data import DataLoader - from monai.config import IgniteInfo -from monai.engines.utils import default_make_latent, default_metric_cmp_fn, default_prepare_batch +from monai.engines.utils import IterationEvents, default_metric_cmp_fn, default_prepare_batch from monai.inferers import Inferer, SimpleInferer from monai.transforms import Transform from monai.utils import GanKeys, min_version, optional_import from monai.utils.enums import CommonKeys, GanKeys +from torch.optim.optimizer import Optimizer +from torch.utils.data import DataLoader if TYPE_CHECKING: from ignite.engine import Engine, EventEnum @@ -31,7 +30,8 @@ Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine") Metric, _ = optional_import("ignite.metrics", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Metric") EventEnum, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "EventEnum") -from monai.engines.trainer import Trainer +from monai.engines.trainer import SupervisedTrainer, Trainer + class VaeGanTrainer(Trainer): """ @@ -199,3 +199,183 @@ def _iteration( GanKeys.GLOSS: g_loss.item(), GanKeys.DLOSS: d_total_loss.item(), } + + +class LDMTrainer(SupervisedTrainer): + """ + Standard supervised training method with image and label, inherits from ``Trainer`` and ``Workflow``. + Args: + device: an object representing the device on which to run. + max_epochs: the total epoch number for trainer to run. + train_data_loader: Ignite engine use data_loader to run, must be Iterable or torch.DataLoader. + network: network to train in the trainer, should be regular PyTorch `torch.nn.Module`. + optimizer: the optimizer associated to the network, should be regular PyTorch optimizer from `torch.optim` + or its subclass. + loss_function: the loss function associated to the optimizer, should be regular PyTorch loss, + which inherit from `torch.nn.modules.loss`. + epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`. + non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously + with respect to the host. For other cases, this argument has no effect. + prepare_batch: function to parse expected data (usually `image`, `label` and other network args) + from `engine.state.batch` for every iteration, for more details please refer to: + https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html. + iteration_update: the callable function for every iteration, expect to accept `engine` + and `engine.state.batch` as inputs, return data will be stored in `engine.state.output`. + if not provided, use `self._iteration()` instead. for more details please refer to: + https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html. + inferer: inference method that execute model forward on input data, like: SlidingWindow, etc. + postprocessing: execute additional transformation for the model output data. + Typically, several Tensor based transforms composed by `Compose`. + key_train_metric: compute metric when every iteration completed, and save average value to + engine.state.metrics when epoch completed. key_train_metric is the main metric to compare and save the + checkpoint into files. + additional_metrics: more Ignite metrics that also attach to Ignite Engine. + metric_cmp_fn: function to compare current key metric with previous best key metric value, + it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update + `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`. + train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like: + CheckpointHandler, StatsHandler, etc. + amp: whether to enable auto-mixed-precision training, default is False. + event_names: additional custom ignite events that will register to the engine. + new events can be a list of str or `ignite.engine.events.EventEnum`. + event_to_attr: a dictionary to map an event to a state attribute, then add to `engine.state`. + for more details, check: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html + #ignite.engine.engine.Engine.register_events. + decollate: whether to decollate the batch-first data to a list of data after model computation, + recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`. + default to `True`. + optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None. + more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html. + to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for + `device`, `non_blocking`. + amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details: + https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast. + """ + + def __init__( + self, + device: str | torch.device, + max_epochs: int, + train_data_loader: Iterable | DataLoader, + network: torch.nn.Module, + autoencoder_model: torch.nn.Module, + optimizer: Optimizer, + loss_function: Callable, + latent_shape: Sequence, + inferer: Inferer, + epoch_length: int | None = None, + non_blocking: bool = False, + prepare_batch: Callable = default_prepare_batch, + iteration_update: Callable[[Engine, Any], Any] | None = None, + postprocessing: Transform | None = None, + key_train_metric: dict[str, Metric] | None = None, + additional_metrics: dict[str, Metric] | None = None, + metric_cmp_fn: Callable = default_metric_cmp_fn, + train_handlers: Sequence | None = None, + amp: bool = False, + event_names: list[str | EventEnum | type[EventEnum]] | None = None, + event_to_attr: dict | None = None, + decollate: bool = True, + optim_set_to_none: bool = False, + to_kwargs: dict | None = None, + amp_kwargs: dict | None = None, + ) -> None: + super().__init__( + device=device, + max_epochs=max_epochs, + train_data_loader=train_data_loader, + network=network, + optimizer=optimizer, + loss_function=loss_function, + inferer=inferer, + optim_set_to_none=optim_set_to_none, + epoch_length=epoch_length, + non_blocking=non_blocking, + prepare_batch=prepare_batch, + iteration_update=iteration_update, + postprocessing=postprocessing, + key_train_metric=key_train_metric, + additional_metrics=additional_metrics, + metric_cmp_fn=metric_cmp_fn, + train_handlers=train_handlers, + amp=amp, + event_names=event_names, + event_to_attr=event_to_attr, + decollate=decollate, + to_kwargs=to_kwargs, + amp_kwargs=amp_kwargs, + ) + + self.latent_shape = latent_shape + self.autoencoder_model = autoencoder_model + + def _iteration(self, engine: LDMTrainer, batchdata: dict[str, torch.Tensor]) -> dict: + """ + Callback function for the Supervised Training processing logic of 1 iteration in Ignite Engine. + Return below items in a dictionary: + - IMAGE: image Tensor data for model input, already moved to device. + - LABEL: label Tensor data corresponding to the image, already moved to device. + - PRED: prediction result of model. + - LOSS: loss value computed by loss function. + Args: + engine: `SupervisedTrainer` to execute operation for an iteration. + batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data. + Raises: + ValueError: When ``batchdata`` is None. + """ + if batchdata is None: + raise ValueError("Must provide batch data for current iteration.") + batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs) + if len(batch) == 2: + images, labels = batch + args: tuple = () + kwargs: dict = {} + else: + images, labels, args, kwargs = batch + # put iteration outputs into engine.state + engine.state.output = {CommonKeys.IMAGE: images} + + # generate noise + noise_shape = [images.shape[0]] + list(self.latent_shape) + noise = torch.randn(noise_shape, dtype=images.dtype).to(images.device) + engine.state.output = {"noise": noise} + + # Create timesteps + timesteps = torch.randint( + 0, engine.inferer.scheduler.num_train_timesteps, (images.shape[0],), device=images.device + ).long() + + def _compute_pred_loss(): + # predicted noise + engine.state.output[CommonKeys.PRED] = engine.inferer( + inputs=images, + autoencoder_model=self.autoencoder_model, + diffusion_model=engine.network, + noise=noise, + timesteps=timesteps, + ) + engine.fire_event(IterationEvents.FORWARD_COMPLETED) + # compute loss + engine.state.output[CommonKeys.LOSS] = engine.loss_function( + engine.state.output[CommonKeys.PRED], noise + ).mean() + engine.fire_event(IterationEvents.LOSS_COMPLETED) + + engine.network.train() + engine.optimizer.zero_grad(set_to_none=engine.optim_set_to_none) + + if engine.amp and engine.scaler is not None: + with torch.cuda.amp.autocast(**engine.amp_kwargs): + _compute_pred_loss() + engine.scaler.scale(engine.state.output[CommonKeys.LOSS]).backward() + engine.fire_event(IterationEvents.BACKWARD_COMPLETED) + engine.scaler.step(engine.optimizer) + engine.scaler.update() + else: + _compute_pred_loss() + engine.state.output[CommonKeys.LOSS].backward() + engine.fire_event(IterationEvents.BACKWARD_COMPLETED) + engine.optimizer.step() + engine.fire_event(IterationEvents.MODEL_COMPLETED) + + return engine.state.output diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 4820d82f..7db25725 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -1,6 +1,5 @@ import torch from generative.losses import PatchAdversarialLoss, PerceptualLoss -from generative.networks.nets import PatchDiscriminator intensity_loss = torch.nn.L1Loss() adv_loss = PatchAdversarialLoss(criterion="least_squares") @@ -26,11 +25,11 @@ def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_percep logits_fake = disc_net(gen_images)[-1] generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False) loss_g = loss_g + adv_weight * generator_loss - + return loss_g - - -def discriminator_loss(gen_images, real_images, disc_net): + + +def discriminator_loss(gen_images, real_images, disc_net): logits_fake = disc_net(gen_images.contiguous().detach())[-1] loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True) logits_real = disc_net(real_images.contiguous().detach())[-1] From 4b1524d7022cfa1dd0a84ba26d7c6ea00358250b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 01:25:43 +0000 Subject: [PATCH 006/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/scripts/losses.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 7db25725..a207ac3b 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -25,11 +25,11 @@ def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_percep logits_fake = disc_net(gen_images)[-1] generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False) loss_g = loss_g + adv_weight * generator_loss - + return loss_g - - -def discriminator_loss(gen_images, real_images, disc_net): + + +def discriminator_loss(gen_images, real_images, disc_net): logits_fake = disc_net(gen_images.contiguous().detach())[-1] loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True) logits_real = disc_net(real_images.contiguous().detach())[-1] From b96fa415804a57ba30b545fa5fa18d0bb1e6c6b3 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 21:19:19 -0700 Subject: [PATCH 007/130] add initi attempt of multi-gpu Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 3 ++- models/brats_mri_generative_diffusion/docs/README.md | 10 +++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index b6fa8d85..07673f3c 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -9,6 +9,7 @@ "ckpt_path": "$@bundle_root + '/models/model_autoencoder.pt'", "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "/workspace/dataset", + "download_brats": true, "train_batch_size": 2, "train_patch_size": [ 112, @@ -139,7 +140,7 @@ "section": "training", "cache_rate": 0, "num_workers": 8, - "download": false, + "download": "@download_brats", "transform": "@train#preprocessing" }, "dataloader": { diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 8d4fbda8..079e948b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. @@ -16,12 +16,16 @@ The provided labelled data was partitioned, based on our own split, into trainin ### Training -Assuming the current directory is the bundle directory, and the dataset was extracted to the directory `./MedNIST`, the following command will train the network for 50 epochs: - +Assuming the current directory is the bundle directory, and the dataset was extracted to the directory `./Task01_BrainTumour`, the following command will train the autoencoder network for 1500 epochs: ``` python -m monai.bundle run --config_file configs/train_autoencoder.json ``` +Or run it with multi-gpu: +``` +torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" +``` + ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` From 1cdc989c3f09df5f88c593b7b49a3709b9328a19 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 04:20:01 +0000 Subject: [PATCH 008/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 079e948b..e2b75921 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. From 46b03ee0c89a71e0b2e886a4a1a8f8dce16f4df3 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 21:27:01 -0700 Subject: [PATCH 009/130] update readme Signed-off-by: Can-Zhao --- .../configs/multi_gpu_train_autoencoder.json | 42 +++++++++++++++++++ .../docs/README.md | 16 ++++--- 2 files changed, 49 insertions(+), 9 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json new file mode 100644 index 00000000..b4cfbece --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json @@ -0,0 +1,42 @@ +{ + "device": "$torch.device(f'cuda:{dist.get_rank()}')", + "gnetwork": { + "_target_": "torch.nn.parallel.DistributedDataParallel", + "module": "$@autoencoder_def.to(@device)", + "device_ids": [ + "@device" + ], + "find_unused_parameters": true + }, + "dnetwork": { + "_target_": "torch.nn.parallel.DistributedDataParallel", + "module": "$@discriminator_def.to(@device)", + "device_ids": [ + "@device" + ], + "find_unused_parameters": true + }, + "train#sampler": { + "_target_": "DistributedSampler", + "dataset": "@train#dataset", + "even_divisible": true, + "shuffle": true + }, + "train#dataloader#sampler": "@train#sampler", + "train#dataloader#shuffle": false, + "train#trainer#train_handlers": "$@train#handlers[: -2 if dist.get_rank() > 0 else None]", + "initialize": [ + "$import torch.distributed as dist", + "$dist.is_initialized() or dist.init_process_group(backend='nccl')", + "$torch.cuda.set_device(@device)", + "$monai.utils.set_determinism(seed=123)", + "$import logging", + "$@train#trainer.logger.setLevel(logging.WARNING if dist.get_rank() > 0 else logging.INFO)" + ], + "run": [ + "$@train#trainer.run()" + ], + "finalize": [ + "$dist.is_initialized() and dist.destroy_process_group()" + ] +} \ No newline at end of file diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 079e948b..f2ba1445 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -1,18 +1,16 @@ # Brats MRI 3D Latent Diffusion Generative Model -This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss, and is based on the [MONAI GAN tutorials](https://github.com/Project-MONAI/tutorials/blob/main/modules/mednist_GAN_workflow_dict.ipynb). - +This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Downloading the Dataset -The training data is from the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018. +The training data is from the [Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018](https://www.med.upenn.edu/sbia/brats2018.html). -Target: 3 tumor subregions -Task: Segmentation +Target: image generatiion +Task: Synthesis Modality: MRI -Size: 285 3D volumes (4 channels each) -The provided labelled data was partitioned, based on our own split, into training (200 studies), validation (42 studies) and testing (43 studies) datasets. +Size: 285 3D volumes (1 channel used) ### Training @@ -40,14 +38,14 @@ python -m monai.bundle run --config_file "['configs/train_autoencoder.json','con The generator can be exported to a Torchscript bundle with the following: ``` -python -m monai.bundle ckpt_export network_def --filepath mednist_gan.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json +python -m monai.bundle ckpt_export autoencoder_def --filepath autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json ``` The model can be loaded without MONAI code after this operation. For example, an image can be generated from a set of random values with: ```python import torch -net = torch.jit.load("mednist_gan.ts") +net = torch.jit.load("autoencoder.ts") latent = torch.rand(1, 64) img = net(latent) # (1,1,64,64) ``` From 5389f178f23447f79f16c9b558ca668ebe15af01 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 04:28:44 +0000 Subject: [PATCH 010/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../configs/multi_gpu_train_autoencoder.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json index b4cfbece..5098ddff 100644 --- a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json @@ -39,4 +39,4 @@ "finalize": [ "$dist.is_initialized() and dist.destroy_process_group()" ] -} \ No newline at end of file +} From 9c4fb68b7fec8d9e2893b9c8dafa809eda3a7591 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 22:10:32 -0700 Subject: [PATCH 011/130] update readme Signed-off-by: Can-Zhao --- .../brats_mri_generative_diffusion/docs/README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 58a3d767..18969790 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -5,6 +5,19 @@ This model is a generator for creating images like the Flair MRIs based on BraTS This is a demonstration network meant to just show the training process for this sort of network with MONAI. +### Install the dependency of MONAI generative models +[MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) +can be installed by +``` +git clone https://github.com/Project-MONAI/GenerativeModels.git +cd GenerativeModels/ +python setup.py install +``` +We also need +``` +pip install lpips +``` + ### Downloading the Dataset The training data is from the [Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018](https://www.med.upenn.edu/sbia/brats2018.html). From 4c1ffbfc27cb7fff3f18acaffb9965fbb688297e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 22:18:36 -0700 Subject: [PATCH 012/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 18969790..9de9c782 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -6,8 +6,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS This is a demonstration network meant to just show the training process for this sort of network with MONAI. ### Install the dependency of MONAI generative models -[MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) -can be installed by +[MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ From f2f74cfd25aea0fd0151be8706e7db4214caf6af Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 22:19:42 -0700 Subject: [PATCH 013/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 9de9c782..2e17a832 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -11,6 +11,7 @@ This is a demonstration network meant to just show the training process for this git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ python setup.py install +cd .. ``` We also need ``` From ff5c944fd56955442debe2b1b97d90aaad3539f1 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 25 Apr 2023 13:41:02 +0800 Subject: [PATCH 014/130] remove unused args Signed-off-by: Yiheng Wang --- .../configs/multi_gpu_train_autoencoder.json | 6 ++---- models/brats_mri_generative_diffusion/scripts/__init__.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json index 5098ddff..e0061ff4 100644 --- a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_autoencoder.json @@ -5,16 +5,14 @@ "module": "$@autoencoder_def.to(@device)", "device_ids": [ "@device" - ], - "find_unused_parameters": true + ] }, "dnetwork": { "_target_": "torch.nn.parallel.DistributedDataParallel", "module": "$@discriminator_def.to(@device)", "device_ids": [ "@device" - ], - "find_unused_parameters": true + ] }, "train#sampler": { "_target_": "DistributedSampler", diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index 85ee5174..71f54e4a 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -10,4 +10,4 @@ # limitations under the License. from . import losses -from . import vae_gan_trainer +from . import ldm_trainer From 2fc81911932f6364fc4f655e4aefdcdd2fb36fc5 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 23:03:34 -0700 Subject: [PATCH 015/130] add StatsHandler Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 5 +++++ .../configs/train_diffusion.json | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 07673f3c..f7d84a8a 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -161,6 +161,11 @@ "save_final": true, "epoch_level": true }, + { + "_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", diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 8e0c2fa5..ac29c79f 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -113,6 +113,11 @@ "save_final": true, "epoch_level": true }, + { + "_target_": "StatsHandler", + "tag_name": "train_diffusion_loss", + "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)" + }, { "_target_": "TensorBoardStatsHandler", "log_dir": "@tf_dir", From 5db0a803acc4531076a362c1aa34256c555f9627 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 24 Apr 2023 23:18:14 -0700 Subject: [PATCH 016/130] change cache Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 2 +- .../configs/train_diffusion.json | 2 +- models/brats_mri_generative_diffusion/scripts/__init__.py | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index f7d84a8a..6d2482b6 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -138,7 +138,7 @@ "root_dir": "@dataset_dir", "task": "Task01_BrainTumour", "section": "training", - "cache_rate": 0, + "cache_rate": 1.0, "num_workers": 8, "download": "@download_brats", "transform": "@train#preprocessing" diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index ac29c79f..960993d9 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -85,7 +85,7 @@ "root_dir": "@dataset_dir", "task": "Task01_BrainTumour", "section": "training", - "cache_rate": 0, + "cache_rate": 1.0, "num_workers": 8, "download": false, "transform": "@train#preprocessing" diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index 71f54e4a..715b0866 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -9,5 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import losses -from . import ldm_trainer +from . import ldm_trainer, losses From ae6641f2ea4e6d86fef54c14134a87bb3587172f Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 09:41:53 -0700 Subject: [PATCH 017/130] update readme Signed-off-by: Can-Zhao --- .../configs/multi_gpu_train_diffusion.json | 40 +++++++++++++++++++ .../configs/train_autoencoder.json | 9 +++-- .../configs/train_diffusion.json | 9 +++-- .../docs/README.md | 15 ++++++- 4 files changed, 64 insertions(+), 9 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json new file mode 100644 index 00000000..e379e7d9 --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json @@ -0,0 +1,40 @@ +{ + "device": "$torch.device(f'cuda:{dist.get_rank()}')", + "autoencoder": { + "_target_": "torch.nn.parallel.DistributedDataParallel", + "module": "$@autoencoder_def.to(@device)", + "device_ids": [ + "@device" + ] + }, + "diffusion": { + "_target_": "torch.nn.parallel.DistributedDataParallel", + "module": "$@diffusion_def.to(@device)", + "device_ids": [ + "@device" + ] + }, + "train#sampler": { + "_target_": "DistributedSampler", + "dataset": "@train#dataset", + "even_divisible": true, + "shuffle": true + }, + "train#dataloader#sampler": "@train#sampler", + "train#dataloader#shuffle": false, + "train#trainer#train_handlers": "$@train#handlers[: -2 if dist.get_rank() > 0 else None]", + "initialize": [ + "$import torch.distributed as dist", + "$dist.is_initialized() or dist.init_process_group(backend='nccl')", + "$torch.cuda.set_device(@device)", + "$monai.utils.set_determinism(seed=123)", + "$import logging", + "$@train#trainer.logger.setLevel(logging.WARNING if dist.get_rank() > 0 else logging.INFO)" + ], + "run": [ + "$@train#trainer.run()" + ], + "finalize": [ + "$dist.is_initialized() and dist.destroy_process_group()" + ] +} diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 6d2482b6..0c2243d1 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -6,10 +6,10 @@ ], "bundle_root": ".", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", - "ckpt_path": "$@bundle_root + '/models/model_autoencoder.pt'", + "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "/workspace/dataset", - "download_brats": true, + "download_brats": false, "train_batch_size": 2, "train_patch_size": [ 112, @@ -153,13 +153,14 @@ "handlers": [ { "_target_": "CheckpointSaver", - "save_dir": "$@bundle_root + '/models'", + "save_dir": "@ckpt_dir", "save_dict": { "model": "@gnetwork" }, "save_interval": 0, "save_final": true, - "epoch_level": true + "epoch_level": true, + "key_metric_filename": "model_autoencoder.pt" }, { "_target_": "StatsHandler", diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 960993d9..b605cedf 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -1,5 +1,5 @@ { - "ckpt_path": "$@bundle_root + '/models/model_ldm.pt'", + "ckpt_dir": "$@bundle_root + '/models'", "train_batch_size": 4, "train_patch_size": [ 144, @@ -87,7 +87,7 @@ "section": "training", "cache_rate": 1.0, "num_workers": 8, - "download": false, + "download": "@download_brats", "transform": "@train#preprocessing" }, "dataloader": { @@ -105,13 +105,14 @@ }, { "_target_": "CheckpointSaver", - "save_dir": "$@bundle_root + '/models'", + "save_dir": "@ckpt_dir", "save_dict": { "model": "@diffusion" }, "save_interval": 0, "save_final": true, - "epoch_level": true + "epoch_level": true, + "key_metric_filename": "model_ldm.pt" }, { "_target_": "StatsHandler", diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 2e17a832..65ff5e2b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -26,9 +26,16 @@ Task: Synthesis Modality: MRI Size: 285 3D volumes (1 channel used) +The dataset can be downloaded automatically at the beggining of training. + ### Training -Assuming the current directory is the bundle directory, and the dataset was extracted to the directory `./Task01_BrainTumour`, the following command will train the autoencoder network for 1500 epochs: +Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs. +If the Brats dataset is not downloaded, run the following command, the data will be downloaded and extracted to `./Task01_BrainTumour`. +``` +python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True +``` +If it is already downloaded, run: ``` python -m monai.bundle run --config_file configs/train_autoencoder.json ``` @@ -37,10 +44,16 @@ Or run it with multi-gpu: ``` torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" ``` +It take 9 hours when trianing with 9 32G GPU. +After the autoencoder is trained, run the following command to train the latent diffusion model. ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` +Or run it with multi-gpu: +``` +torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_diffusion.json']" +``` ### Inference From c906e5fe532ff269fe64e66c44cd8e54b554b229 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 13:33:54 -0700 Subject: [PATCH 018/130] add inference Signed-off-by: Can-Zhao --- .../configs/inference.json | 44 +++++++++++++ .../configs/multi_gpu_train_diffusion.json | 13 ++-- .../configs/train_autoencoder.json | 4 +- .../configs/train_diffusion.json | 10 ++- .../docs/README.md | 9 ++- .../scripts/__init__.py | 2 +- .../scripts/ldm_sampler.py | 63 +++++++++++++++++++ 7 files changed, 130 insertions(+), 15 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/configs/inference.json create mode 100644 models/brats_mri_generative_diffusion/scripts/ldm_sampler.py diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json new file mode 100644 index 00000000..8f65f472 --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -0,0 +1,44 @@ +{ + "imports": [ + "$import torch", + "$from datetime import datetime", + "$from pathlib import Path" + ], + "bundle_root": ".", + "model_dir": "$@bundle_root + '/models'", + "output_dir": "$@bundle_root + '/output'", + "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", + "load_diffusion_path": "$@model_dir + '/model_ldm.pt'", + "load_diffusion": "$@diffusion_def.load_state_dict(torch.load(@load_diffusion_path))", + "diffusion": "$@diffusion_def.to(@device)", + "noise_scheduler": { + "_target_": "generative.networks.schedulers.DDIMScheduler", + "_requires_": [ + "@load_diffusion", + "@load_autoencoder" + ], + "num_train_timesteps": 1000, + "beta_start": 0.0015, + "beta_end": 0.0195, + "beta_schedule": "scaled_linear", + "clip_sample": false + }, + "noise": "$torch.randn([1]+@latent_shape).to(@device)", + "set_timesteps": "$@noise_scheduler.set_timesteps(num_inference_steps=50)", + "sampler": { + "_target_": "scripts.ldm_sampler.LDMSampler", + "_requires_": "@set_timesteps" + }, + "sample": "$@sampler.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", + "saver": { + "_target_": "SaveImage", + "_requires_": "@create_output_dir", + "output_dir": "@output_dir", + "output_postfix": "@output_postfix" + }, + "run": [ + "$@saver(@sample)" + ] +} \ No newline at end of file diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json index e379e7d9..c2e7116d 100644 --- a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json @@ -1,18 +1,12 @@ { "device": "$torch.device(f'cuda:{dist.get_rank()}')", - "autoencoder": { - "_target_": "torch.nn.parallel.DistributedDataParallel", - "module": "$@autoencoder_def.to(@device)", - "device_ids": [ - "@device" - ] - }, "diffusion": { "_target_": "torch.nn.parallel.DistributedDataParallel", "module": "$@diffusion_def.to(@device)", "device_ids": [ "@device" - ] + ], + "find_unused_parameters": true }, "train#sampler": { "_target_": "DistributedSampler", @@ -32,6 +26,9 @@ "$@train#trainer.logger.setLevel(logging.WARNING if dist.get_rank() > 0 else logging.INFO)" ], "run": [ + "@load_autoencoder", + "$@autoencoder.eval()", + "$print('scale factor:',@scale_factor)", "$@train#trainer.run()" ], "finalize": [ diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 0c2243d1..350431b7 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -8,7 +8,7 @@ "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/dataset", + "dataset_dir": "$@bundle_root", "download_brats": false, "train_batch_size": 2, "train_patch_size": [ @@ -160,7 +160,7 @@ "save_interval": 0, "save_final": true, "epoch_level": true, - "key_metric_filename": "model_autoencoder.pt" + "final_filename": "model_autoencoder.pt" }, { "_target_": "StatsHandler", diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index b605cedf..83652098 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -52,9 +52,12 @@ ], "gamma": 0.1 }, - "scale_factor": 1.0, + "scale_factor": "$1.0/torch.std(@autoencoder.to(@device).encode_stage_2_inputs(monai.utils.first(@train#dataloader)['image'].to(@device))).item()", "noise_scheduler": { "_target_": "generative.networks.schedulers.DDPMScheduler", + "_requires_": [ + "@load_autoencoder" + ], "beta_schedule": "scaled_linear", "num_train_timesteps": 1000, "beta_start": 0.0015, @@ -112,7 +115,7 @@ "save_interval": 0, "save_final": true, "epoch_level": true, - "key_metric_filename": "model_ldm.pt" + "final_filename": "model_ldm.pt" }, { "_target_": "StatsHandler", @@ -145,6 +148,9 @@ "$monai.utils.set_determinism(seed=0)" ], "run": [ + "@load_autoencoder", + "$@autoencoder.eval()", + "$print('scale factor:',@scale_factor)", "$@train#trainer.run()" ] } diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 65ff5e2b..b0f755e7 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -35,7 +35,7 @@ If the Brats dataset is not downloaded, run the following command, the data will ``` python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True ``` -If it is already downloaded, run: +If it is already downloaded, check if `"dataset_dir"` in `configs/train_autoencoder.json`has folder `Task01_BrainTumour`. If so, run: ``` python -m monai.bundle run --config_file configs/train_autoencoder.json ``` @@ -44,12 +44,14 @@ Or run it with multi-gpu: ``` torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" ``` -It take 9 hours when trianing with 9 32G GPU. +It take 9 hours when training with 9 32G GPU. After the autoencoder is trained, run the following command to train the latent diffusion model. ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` +It will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. + Or run it with multi-gpu: ``` torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_diffusion.json']" @@ -57,6 +59,9 @@ torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config ### Inference +``` +python -m monai.bundle save_nii --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" +``` diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index 715b0866..ac5b5062 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import ldm_trainer, losses +from . import ldm_trainer, losses, ldm_sampler diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py new file mode 100644 index 00000000..32d213d3 --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py @@ -0,0 +1,63 @@ +# 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. + +from __future__ import annotations + +import torch +import torch.nn as nn +from monai.utils import optional_import +from torch.cuda.amp import autocast + +tqdm, has_tqdm = optional_import("tqdm", name="tqdm") + + +class LDMSampler: + def __init__(self) -> None: + super().__init__() + + @torch.no_grad() + def sampling_fn( + self, + input_noise: torch.Tensor, + autoencoder_model: nn.Module, + diffusion_model: nn.Module, + scheduler: nn.Module, + conditioning: torch.Tensor|None=None, + ) -> torch.Tensor: + if has_tqdm: + progress_bar = tqdm(scheduler.timesteps) + else: + progress_bar = iter(scheduler.timesteps) + + image = input_noise + if conditioning is not None: + cond_concat = conditioning.squeeze(1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + cond_concat = cond_concat.expand(list(cond_concat.shape[0:2]) + list(input_noise.shape[2:])) + + + for t in progress_bar: + with torch.no_grad(): + if conditioning is not None: + input_t = torch.cat((image, cond_concat), dim=1) + else: + input_t = image + model_output = diffusion_model( + input_t, + timesteps=torch.Tensor((t,)).to(input_noise.device).long(), + context=conditioning, + ) + image, _ = scheduler.step(model_output, t, image) + + with torch.no_grad(): + with autocast(): + sample = autoencoder_model.decode_stage_2_outputs(image) + + return sample \ No newline at end of file From b08e8fd72efd09acdb7b0269229102b96645ba3b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 20:34:32 +0000 Subject: [PATCH 019/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/configs/inference.json | 4 ++-- models/brats_mri_generative_diffusion/scripts/ldm_sampler.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 8f65f472..9ae50e4a 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -38,7 +38,7 @@ "output_dir": "@output_dir", "output_postfix": "@output_postfix" }, - "run": [ + "run": [ "$@saver(@sample)" ] -} \ No newline at end of file +} diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py index 32d213d3..3bf9bf9e 100644 --- a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py @@ -41,7 +41,7 @@ def sampling_fn( if conditioning is not None: cond_concat = conditioning.squeeze(1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) cond_concat = cond_concat.expand(list(cond_concat.shape[0:2]) + list(input_noise.shape[2:])) - + for t in progress_bar: with torch.no_grad(): @@ -60,4 +60,4 @@ def sampling_fn( with autocast(): sample = autoencoder_model.decode_stage_2_outputs(image) - return sample \ No newline at end of file + return sample From e759e0065429cfd81a77963567ef69d09ebd7a08 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 13:42:06 -0700 Subject: [PATCH 020/130] update meta Signed-off-by: Can-Zhao --- .../brats_mri_generative_diffusion/configs/metadata.json | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 3a7f2ccf..c4855944 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -4,11 +4,11 @@ "changelog": { "1.0.0": "Initial release" }, - "monai_version": "1.1.0", - "pytorch_version": "1.13.0", - "numpy_version": "1.22.4", + "monai_version": "1.2.0", + "pytorch_version": "2.0.0", + "numpy_version": "1.22.2", "optional_packages_version": { - "nibabel": "4.0.1", + "nibabel": "5.1.0", "generative": "0.1.0" }, "task": "BraTS MRI image synthesis", From 4786c71461c58c600ecb5a3b46259adf637c9f78 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 13:59:19 -0700 Subject: [PATCH 021/130] update readme Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 3 ++- .../configs/train_diffusion.json | 3 ++- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 350431b7..f03f9486 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -171,7 +171,8 @@ "_target_": "TensorBoardStatsHandler", "log_dir": "@tf_dir", "tag_name": "train_loss", - "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]" + "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]", + "iteration_log":false } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 83652098..8921dbe8 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -126,7 +126,8 @@ "_target_": "TensorBoardStatsHandler", "log_dir": "@tf_dir", "tag_name": "train_diffusion_loss", - "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)" + "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)", + "iteration_log":false } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index b0f755e7..2b0940e7 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -54,13 +54,13 @@ It will print out the scale factor of the latent feature space. If your autoenco Or run it with multi-gpu: ``` -torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_diffusion.json']" +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_diffusion.json']" ``` ### Inference ``` -python -m monai.bundle save_nii --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" +python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` From 47eb1b48ad56bc385ddbf8cb07c275803e6f34ea Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 20:59:32 +0000 Subject: [PATCH 022/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../configs/train_autoencoder.json | 2 +- .../brats_mri_generative_diffusion/configs/train_diffusion.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index f03f9486..59d0d598 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -172,7 +172,7 @@ "log_dir": "@tf_dir", "tag_name": "train_loss", "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]", - "iteration_log":false + "iteration_log": false } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 8921dbe8..a5c056c1 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -127,7 +127,7 @@ "log_dir": "@tf_dir", "tag_name": "train_diffusion_loss", "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)", - "iteration_log":false + "iteration_log": false } ], "trainer": { From 8ae26383e783ca46890842b3a2051765ddeb6bc8 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:22:05 -0700 Subject: [PATCH 023/130] corrent inference Signed-off-by: Can-Zhao --- .../configs/inference.json | 7 ++++--- .../configs/train_diffusion.json | 4 ++-- .../scripts/__init__.py | 2 +- .../scripts/utils.py | 18 ++++++++++++++++++ 4 files changed, 25 insertions(+), 6 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/scripts/utils.py diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 9ae50e4a..bc36351e 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -38,7 +38,8 @@ "output_dir": "@output_dir", "output_postfix": "@output_postfix" }, - "run": [ - "$@saver(@sample)" + "generated_image": "$@sample", + "run": [ + "$@saver(@generated_image[0])" ] -} +} \ No newline at end of file diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index a5c056c1..989658e7 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -52,7 +52,7 @@ ], "gamma": 0.1 }, - "scale_factor": "$1.0/torch.std(@autoencoder.to(@device).encode_stage_2_inputs(monai.utils.first(@train#dataloader)['image'].to(@device))).item()", + "scale_factor": "$scripts.utils.compute_scale_factor(@autoencoder,@train#dataloader,@device)", "noise_scheduler": { "_target_": "generative.networks.schedulers.DDPMScheduler", "_requires_": [ @@ -127,7 +127,7 @@ "log_dir": "@tf_dir", "tag_name": "train_diffusion_loss", "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)", - "iteration_log": false + "iteration_log":false } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index ac5b5062..d7b349ba 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import ldm_trainer, losses, ldm_sampler +from . import ldm_trainer, losses, ldm_sampler, utils diff --git a/models/brats_mri_generative_diffusion/scripts/utils.py b/models/brats_mri_generative_diffusion/scripts/utils.py new file mode 100644 index 00000000..4b43d743 --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/utils.py @@ -0,0 +1,18 @@ +# 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 + +import torch +import monai +def compute_scale_factor(autoencoder,train_loader,device): + with torch.no_grad(): + check_data = monai.utils.first(train_loader) + z = autoencoder.encode_stage_2_inputs(check_data["image"].to(device)) + scale_factor = 1 / torch.std(z) + return scale_factor.item() \ No newline at end of file From 67be3f394f5914c43641e56ca20793680978bf42 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 23:22:19 +0000 Subject: [PATCH 024/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/configs/inference.json | 4 ++-- .../configs/train_diffusion.json | 2 +- models/brats_mri_generative_diffusion/scripts/utils.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index bc36351e..8463c157 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -39,7 +39,7 @@ "output_postfix": "@output_postfix" }, "generated_image": "$@sample", - "run": [ + "run": [ "$@saver(@generated_image[0])" ] -} \ No newline at end of file +} diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 989658e7..fe7cb7d3 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -127,7 +127,7 @@ "log_dir": "@tf_dir", "tag_name": "train_diffusion_loss", "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)", - "iteration_log":false + "iteration_log": false } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/scripts/utils.py b/models/brats_mri_generative_diffusion/scripts/utils.py index 4b43d743..dbf51924 100644 --- a/models/brats_mri_generative_diffusion/scripts/utils.py +++ b/models/brats_mri_generative_diffusion/scripts/utils.py @@ -15,4 +15,4 @@ def compute_scale_factor(autoencoder,train_loader,device): check_data = monai.utils.first(train_loader) z = autoencoder.encode_stage_2_inputs(check_data["image"].to(device)) scale_factor = 1 / torch.std(z) - return scale_factor.item() \ No newline at end of file + return scale_factor.item() From 2f7c48de058714f6f197c0e18c7aebf6ba1158b8 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:22:55 -0700 Subject: [PATCH 025/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/scripts/utils.py b/models/brats_mri_generative_diffusion/scripts/utils.py index 4b43d743..41f8a7e1 100644 --- a/models/brats_mri_generative_diffusion/scripts/utils.py +++ b/models/brats_mri_generative_diffusion/scripts/utils.py @@ -15,4 +15,5 @@ def compute_scale_factor(autoencoder,train_loader,device): check_data = monai.utils.first(train_loader) z = autoencoder.encode_stage_2_inputs(check_data["image"].to(device)) scale_factor = 1 / torch.std(z) - return scale_factor.item() \ No newline at end of file + return scale_factor.item() + \ No newline at end of file From 66d64954da1bac0d4744e48fb93e97f86212f7d9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 25 Apr 2023 23:25:04 +0000 Subject: [PATCH 026/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/scripts/ldm_sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py index 32d213d3..3bf9bf9e 100644 --- a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py @@ -41,7 +41,7 @@ def sampling_fn( if conditioning is not None: cond_concat = conditioning.squeeze(1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) cond_concat = cond_concat.expand(list(cond_concat.shape[0:2]) + list(input_noise.shape[2:])) - + for t in progress_bar: with torch.no_grad(): @@ -60,4 +60,4 @@ def sampling_fn( with autocast(): sample = autoencoder_model.decode_stage_2_outputs(image) - return sample \ No newline at end of file + return sample From e684ea393f4cfef1f887f31d46cddc4695d1b7f9 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:32:36 -0700 Subject: [PATCH 027/130] reformat Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 3 +-- .../configs/train_diffusion.json | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 59d0d598..350431b7 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -171,8 +171,7 @@ "_target_": "TensorBoardStatsHandler", "log_dir": "@tf_dir", "tag_name": "train_loss", - "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]", - "iteration_log": false + "output_transform": "$lambda x: monai.handlers.from_engine(['g_loss'], first=True)(x)[0]" } ], "trainer": { diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index fe7cb7d3..32fae5a7 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -126,8 +126,7 @@ "_target_": "TensorBoardStatsHandler", "log_dir": "@tf_dir", "tag_name": "train_diffusion_loss", - "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)", - "iteration_log": false + "output_transform": "$lambda x: monai.handlers.from_engine(['loss'], first=True)(x)" } ], "trainer": { From 1ab7f9a02fcd1632adfc01946dc95f15d0940203 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:33:28 -0700 Subject: [PATCH 028/130] update meta Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/metadata.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index c4855944..d97266a5 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -4,7 +4,7 @@ "changelog": { "1.0.0": "Initial release" }, - "monai_version": "1.2.0", + "monai_version": "1.2.0rc5", "pytorch_version": "2.0.0", "numpy_version": "1.22.2", "optional_packages_version": { From af4ac4c4e9993689bcc13252c7b9a81f6a6a0c28 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:53:02 -0700 Subject: [PATCH 029/130] update readme Signed-off-by: Can-Zhao --- .../docs/README.md | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 2b0940e7..fc356855 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -62,24 +62,36 @@ torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` - +The generated image will be saved to `./output/0` ### Export -The generator can be exported to a Torchscript bundle with the following: +The autoencoder and latent diffusion generators can be exported to a Torchscript bundle with the following: ``` -python -m monai.bundle ckpt_export autoencoder_def --filepath autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json +python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" +python -m monai.bundle ckpt_export diffusion_def --filepath models/model_ldm.ts --ckpt_file models/model_ldm.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` -The model can be loaded without MONAI code after this operation. For example, an image can be generated from a set of random values with: +The models can be loaded after this operation. For example, an image can be generated from a set of random values with: ```python import torch -net = torch.jit.load("autoencoder.ts") -latent = torch.rand(1, 64) -img = net(latent) # (1,1,64,64) +from scripts.ldm_sampler import LDMSampler +from generative.networks.schedulers import DDIMScheduler +autoencoder = torch.jit.load("models/model_autoencoder.ts") +diffusion = torch.jit.load("models/model_ldm.ts") +noise = torch.rand(1, 8,36,44,28).to(device) +noise_scheduler = DDIMScheduler( + num_train_timesteps=1000, + beta_start= 0.0015, + beta_end=0.0195, + beta_schedule= "scaled_linear", + clip_sample= False) +noise_scheduler.set_timesteps(num_inference_steps=50) +sampler = LDMSampler() +img = sampler.sampling_fn(noise, autoencoder, diffusion, noise_scheduler) ``` # License From bb5fab495fa316931f6587231f3baf2e46e77f32 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 16:56:20 -0700 Subject: [PATCH 030/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/ldm_sampler.py | 3 +-- models/brats_mri_generative_diffusion/scripts/losses.py | 2 ++ models/brats_mri_generative_diffusion/scripts/utils.py | 4 +++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py index 3bf9bf9e..d39d23e4 100644 --- a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py @@ -30,7 +30,7 @@ def sampling_fn( autoencoder_model: nn.Module, diffusion_model: nn.Module, scheduler: nn.Module, - conditioning: torch.Tensor|None=None, + conditioning: torch.Tensor | None = None, ) -> torch.Tensor: if has_tqdm: progress_bar = tqdm(scheduler.timesteps) @@ -42,7 +42,6 @@ def sampling_fn( cond_concat = conditioning.squeeze(1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) cond_concat = cond_concat.expand(list(cond_concat.shape[0:2]) + list(input_noise.shape[2:])) - for t in progress_bar: with torch.no_grad(): if conditioning is not None: diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index dfdf88dc..a075d1fa 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -22,10 +22,12 @@ # If too small, latent space will not be regularized enough for the diffusion model kl_weight = 1e-7 + def KL_loss(z_mu, z_sigma): kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) return torch.sum(kl_loss) / kl_loss.shape[0] + def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) kl_loss = KL_loss(z_mu, z_sigma) diff --git a/models/brats_mri_generative_diffusion/scripts/utils.py b/models/brats_mri_generative_diffusion/scripts/utils.py index dbf51924..0cdef7d5 100644 --- a/models/brats_mri_generative_diffusion/scripts/utils.py +++ b/models/brats_mri_generative_diffusion/scripts/utils.py @@ -10,7 +10,9 @@ import torch import monai -def compute_scale_factor(autoencoder,train_loader,device): + + +def compute_scale_factor(autoencoder, train_loader, device): with torch.no_grad(): check_data = monai.utils.first(train_loader) z = autoencoder.encode_stage_2_inputs(check_data["image"].to(device)) From f2b7b6ccb92447c7218e7fec3ef58565040e203b Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:01:19 -0700 Subject: [PATCH 031/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/__init__.py | 2 +- models/brats_mri_generative_diffusion/scripts/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/__init__.py b/models/brats_mri_generative_diffusion/scripts/__init__.py index d7b349ba..2041a809 100644 --- a/models/brats_mri_generative_diffusion/scripts/__init__.py +++ b/models/brats_mri_generative_diffusion/scripts/__init__.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import ldm_trainer, losses, ldm_sampler, utils +from . import ldm_sampler, ldm_trainer, losses, utils diff --git a/models/brats_mri_generative_diffusion/scripts/utils.py b/models/brats_mri_generative_diffusion/scripts/utils.py index 0cdef7d5..57a6744e 100644 --- a/models/brats_mri_generative_diffusion/scripts/utils.py +++ b/models/brats_mri_generative_diffusion/scripts/utils.py @@ -8,8 +8,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and -import torch import monai +import torch def compute_scale_factor(autoencoder, train_loader, device): From 8fa54052b383d0ecf910b2a802822f775f713b72 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:11:03 -0700 Subject: [PATCH 032/130] refactor Signed-off-by: Can-Zhao --- .../configs/multi_gpu_train_diffusion.json | 21 ------------------- .../docs/README.md | 2 +- 2 files changed, 1 insertion(+), 22 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json index c2e7116d..89a381a7 100644 --- a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json @@ -1,5 +1,4 @@ { - "device": "$torch.device(f'cuda:{dist.get_rank()}')", "diffusion": { "_target_": "torch.nn.parallel.DistributedDataParallel", "module": "$@diffusion_def.to(@device)", @@ -8,30 +7,10 @@ ], "find_unused_parameters": true }, - "train#sampler": { - "_target_": "DistributedSampler", - "dataset": "@train#dataset", - "even_divisible": true, - "shuffle": true - }, - "train#dataloader#sampler": "@train#sampler", - "train#dataloader#shuffle": false, - "train#trainer#train_handlers": "$@train#handlers[: -2 if dist.get_rank() > 0 else None]", - "initialize": [ - "$import torch.distributed as dist", - "$dist.is_initialized() or dist.init_process_group(backend='nccl')", - "$torch.cuda.set_device(@device)", - "$monai.utils.set_determinism(seed=123)", - "$import logging", - "$@train#trainer.logger.setLevel(logging.WARNING if dist.get_rank() > 0 else logging.INFO)" - ], "run": [ "@load_autoencoder", "$@autoencoder.eval()", "$print('scale factor:',@scale_factor)", "$@train#trainer.run()" - ], - "finalize": [ - "$dist.is_initialized() and dist.destroy_process_group()" ] } diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index fc356855..cd415aa8 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -54,7 +54,7 @@ It will print out the scale factor of the latent feature space. If your autoenco Or run it with multi-gpu: ``` -torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_diffusion.json']" +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" ``` From 4484e1e8e18d99082fe70f159c971800cce789ab Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:24:01 -0700 Subject: [PATCH 033/130] update readme Signed-off-by: Can-Zhao --- .../docs/README.md | 46 +++++++++++++++++-- 1 file changed, 42 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index cd415aa8..476a3166 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -5,7 +5,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS This is a demonstration network meant to just show the training process for this sort of network with MONAI. -### Install the dependency of MONAI generative models +## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` git clone https://github.com/Project-MONAI/GenerativeModels.git @@ -18,7 +18,7 @@ We also need pip install lpips ``` -### Downloading the Dataset +## Downloading the Dataset The training data is from the [Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018](https://www.med.upenn.edu/sbia/brats2018.html). Target: image generatiion @@ -28,7 +28,45 @@ Size: 285 3D volumes (1 channel used) The dataset can be downloaded automatically at the beggining of training. -### Training +## Training configuration +If user has GPU memory smaller than 32G, then please decrease the `"train_batch_size"` in `configs/train_autoencoder.json` and `configs/train_diffusion.json`. + +### Training configuration of autoencoder +The training of autoencoder was performed with the following: + +- GPU: at least 32GB GPU memory +- Actual Model Input: 112 x 128 x 80 +- AMP: False +- Optimizer: Adam +- Learning Rate: 1e-5 +- Loss: L1 loss, perceptual loss, adversianl loss, GAN BCE loss + +#### Input +1 channel 3D MRI patches + +#### Output +- 1 channel 3D MRI reconstructed patches +- 8 channel mean of latent features +- 8 channel standard deviation of latent features + +### Training configuration of difusion model +The training of latent diffusion model was performed with the following: + +- GPU: at least 32GB GPU memory +- Actual Model Input: 144 x 176 x 112 +- AMP: False +- Optimizer: Adam +- Learning Rate: 1e-5 +- Loss: MSE loss + +#### Input +8 channel random noise + +#### Output +- 1 channel 3D MRI reconstructed images + + +## MONAI Bundle Commands Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs. If the Brats dataset is not downloaded, run the following command, the data will be downloaded and extracted to `./Task01_BrainTumour`. @@ -44,7 +82,7 @@ Or run it with multi-gpu: ``` torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" ``` -It take 9 hours when training with 9 32G GPU. +It take 9 hours when training with 8 32G GPU. After the autoencoder is trained, run the following command to train the latent diffusion model. ``` From 5732c8f896beafb697ff51ff41c39c57f39f92a9 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:26:25 -0700 Subject: [PATCH 034/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 476a3166..99766801 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -53,7 +53,7 @@ The training of autoencoder was performed with the following: The training of latent diffusion model was performed with the following: - GPU: at least 32GB GPU memory -- Actual Model Input: 144 x 176 x 112 +- Actual Model Input: 36 x 44 x 28 - AMP: False - Optimizer: Adam - Learning Rate: 1e-5 From 1237e88582bd0dda924e19b9e99be5049fded308 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:27:16 -0700 Subject: [PATCH 035/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 99766801..4f5a7ec7 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -60,7 +60,7 @@ The training of latent diffusion model was performed with the following: - Loss: MSE loss #### Input -8 channel random noise +8 channel random noise which represents latent features #### Output - 1 channel 3D MRI reconstructed images From 28a0c28894e0bfef144755f3df32b6a40dd02576 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:28:06 -0700 Subject: [PATCH 036/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 4f5a7ec7..fb370a36 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -39,7 +39,7 @@ The training of autoencoder was performed with the following: - AMP: False - Optimizer: Adam - Learning Rate: 1e-5 -- Loss: L1 loss, perceptual loss, adversianl loss, GAN BCE loss +- Loss: L1 loss, perceptual loss, KL divergence loss, adversianl loss, GAN BCE loss #### Input 1 channel 3D MRI patches From a9ad5d707af2e0cfae545e0217a5ac3195d90feb Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:35:45 -0700 Subject: [PATCH 037/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index fb370a36..5d761625 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -1,7 +1,11 @@ +# Model Overview +A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative Model. -# Brats MRI 3D Latent Diffusion Generative Model +This model is trained on BraTS 2018 data (https://luna16.grand-challenge.org/Home/), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). -This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data.. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. +![model workflow placeholder, will update](https://www.google.com/url?sa=i&url=https%3A%2F%2Ftowardsdatascience.com%2Fpaper-explained-high-resolution-image-synthesis-with-latent-diffusion-models-f372f7636d42&psig=AOvVaw2X9URb15ch3IEeql1vDRRO&ust=1682555678960000&source=images&cd=vfe&ved=0CBAQjRxqFwoTCOCzmr-mxv4CFQAAAAAdAAAAABAE) + +This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. This is a demonstration network meant to just show the training process for this sort of network with MONAI. @@ -60,10 +64,10 @@ The training of latent diffusion model was performed with the following: - Loss: MSE loss #### Input -8 channel random noise which represents latent features +8 channel random noise #### Output -- 1 channel 3D MRI reconstructed images +- 8 channel predicted added noise ## MONAI Bundle Commands From 954ee364cdc7643d541d01f1c8104dc28223332b Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:36:41 -0700 Subject: [PATCH 038/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 5d761625..a2f9c847 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -67,7 +67,7 @@ The training of latent diffusion model was performed with the following: 8 channel random noise #### Output -- 8 channel predicted added noise +8 channel predicted added noise ## MONAI Bundle Commands From 662d2dfb0c6d19a3a0507c8444d10bde1a3da961 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:37:59 -0700 Subject: [PATCH 039/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index a2f9c847..b1a06391 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2018 data (https://luna16.grand-challenge.org/Hom This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -This is a demonstration network meant to just show the training process for this sort of network with MONAI. +This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by From 3cff16ef8df3f94534ef4601a2f63312d9cc89c5 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 17:38:29 -0700 Subject: [PATCH 040/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index b1a06391..984b528f 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative This model is trained on BraTS 2018 data (https://luna16.grand-challenge.org/Home/), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). -![model workflow placeholder, will update](https://www.google.com/url?sa=i&url=https%3A%2F%2Ftowardsdatascience.com%2Fpaper-explained-high-resolution-image-synthesis-with-latent-diffusion-models-f372f7636d42&psig=AOvVaw2X9URb15ch3IEeql1vDRRO&ust=1682555678960000&source=images&cd=vfe&ved=0CBAQjRxqFwoTCOCzmr-mxv4CFQAAAAAdAAAAABAE) +![model workflow placeholder, will update](https://miro.medium.com/v2/resize:fit:1256/1*3Jjlrb-TB8hPpIexoCKpZQ.png) This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. From 39ed1047d44c69b39742190f941ffa70e7ae0b53 Mon Sep 17 00:00:00 2001 From: monai-bot Date: Wed, 26 Apr 2023 02:57:17 +0000 Subject: [PATCH 041/130] [MONAI] code formatting Signed-off-by: monai-bot --- models/brats_mri_generative_diffusion/scripts/ldm_sampler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py index d39d23e4..392d3333 100644 --- a/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_sampler.py @@ -49,9 +49,7 @@ def sampling_fn( else: input_t = image model_output = diffusion_model( - input_t, - timesteps=torch.Tensor((t,)).to(input_noise.device).long(), - context=conditioning, + input_t, timesteps=torch.Tensor((t,)).to(input_noise.device).long(), context=conditioning ) image, _ = scheduler.step(model_output, t, image) From f93079b3e22a6c3e86553992fd8db54715813613 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 20:23:10 -0700 Subject: [PATCH 042/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/losses.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index a075d1fa..5a04574f 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -23,14 +23,14 @@ kl_weight = 1e-7 -def KL_loss(z_mu, z_sigma): +def kl_loss(z_mu, z_sigma): kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) return torch.sum(kl_loss) / kl_loss.shape[0] def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) - kl_loss = KL_loss(z_mu, z_sigma) + kl_loss = kl_loss(z_mu, z_sigma) p_loss = loss_perceptual(gen_images.float(), real_images.float()) loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss From 3baa9ff0bfdd0768602bbf19c76c5e4de8c9aa05 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 20:26:27 -0700 Subject: [PATCH 043/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/ldm_trainer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py b/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py index 16367671..c1a21bfa 100644 --- a/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py +++ b/models/brats_mri_generative_diffusion/scripts/ldm_trainer.py @@ -11,14 +11,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Callable, Sequence +from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence import torch from monai.config import IgniteInfo from monai.engines.utils import IterationEvents, default_metric_cmp_fn, default_prepare_batch from monai.inferers import Inferer, SimpleInferer from monai.transforms import Transform -from monai.utils import GanKeys, min_version, optional_import +from monai.utils import min_version, optional_import from monai.utils.enums import CommonKeys, GanKeys from torch.optim.optimizer import Optimizer from torch.utils.data import DataLoader @@ -173,7 +173,6 @@ def _iteration( raise ValueError("must provide batch data for current iteration.") d_input = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)[0] - batch_size = engine.data_loader.batch_size # type: ignore g_input = d_input g_output, z_mu, z_sigma = engine.g_inferer(g_input, engine.g_network) From 1bae63fd92266032fefaee2d372237576d240902 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 20:30:59 -0700 Subject: [PATCH 044/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/losses.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 5a04574f..9be021bf 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -24,8 +24,8 @@ def kl_loss(z_mu, z_sigma): - kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) - return torch.sum(kl_loss) / kl_loss.shape[0] + kl_loss_all = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) + return torch.sum(kl_loss_all) / kl_loss_all.shape[0] def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): From c721cedb936138e3f88e16619ac870e011fb2a17 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 20:32:19 -0700 Subject: [PATCH 045/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/losses.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 9be021bf..8a0deec5 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -23,14 +23,14 @@ kl_weight = 1e-7 -def kl_loss(z_mu, z_sigma): - kl_loss_all = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) - return torch.sum(kl_loss_all) / kl_loss_all.shape[0] +def compute_kl_loss(z_mu, z_sigma): + kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) + return torch.sum(kl_loss) / kl_loss.shape[0] def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) - kl_loss = kl_loss(z_mu, z_sigma) + kl_loss = compute_kl_loss(z_mu, z_sigma) p_loss = loss_perceptual(gen_images.float(), real_images.float()) loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss From ca54b30f9c41bbe031e98ea211390d9abcfdbabb Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 25 Apr 2023 20:40:43 -0700 Subject: [PATCH 046/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 984b528f..60185215 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -63,12 +63,18 @@ The training of latent diffusion model was performed with the following: - Learning Rate: 1e-5 - Loss: MSE loss -#### Input -8 channel random noise +#### Training Input +8 channel noisy latent features -#### Output +#### Training Output 8 channel predicted added noise +#### Inference Input +8 channel noise + +#### Inference Output +8 channel denoised latent features + ## MONAI Bundle Commands From 91fc60ca870e36a8ad2faf8270c37e560eb2a681 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Wed, 26 Apr 2023 11:56:15 +0800 Subject: [PATCH 047/130] add install script for ci test Signed-off-by: Yiheng Wang --- CONTRIBUTING.md | 5 +++++ ci/bundle_custom_data.py | 7 +++++++ ...l_brats_mri_generative_diffusion_dependency.sh | 15 +++++++++++++++ ci/verify_bundle.py | 8 ++++++++ .../configs/metadata.json | 2 +- 5 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 4405a4b5..f9a105a0 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -84,6 +84,11 @@ In train config file (if exists), please follow the following requirements in or We prepared several premerge CI tests to verify your bundle. +### Prepare dependencies + +For dependencies that support `pip install` command, please put them into `optional_packages_version` in `configs/metadata.json` ([click here for instance](https://github.com/Project-MONAI/model-zoo/blob/dev/models/brats_mri_segmentation/configs/metadata.json)), and the CI test program will extract and install all libraries directly before running tests. +For dependencies that require multiple steps to install, please prepare a install script in `ci/install_scripts`, and put the bundle name and the script path into `install_dependency_dict` in `ci/bundle_custom_data.py` ([click here for instance](https://github.com/Project-MONAI/model-zoo/tree/dev/ci/install_scripts/)). + ### Necessary verifications 1. Check if necessary files are existing in your bundle. diff --git a/ci/bundle_custom_data.py b/ci/bundle_custom_data.py index 4fc9e05b..7bb2e4d2 100644 --- a/ci/bundle_custom_data.py +++ b/ci/bundle_custom_data.py @@ -30,6 +30,13 @@ "mednist_reg", ] +# This dict is used for our CI tests to install required dependencies that cannot be installed by `pip install` directly. +# If a bundle has this kind of dependencies, please add the bundle name (key), and the path of the install script (value) +# into the dict. +install_dependency_dict = { + "brats_mri_generative_diffusion": "ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh" +} + # This list is used for our CI tests to determine whether a bundle supports TensorRT export. Related # test will be employed for bundles in the list. include_verify_tensorrt_list = ["spleen_ct_segmentation", "endoscopic_tool_segmentation", "pathology_tumor_detection"] diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh new file mode 100644 index 00000000..08dd23e8 --- /dev/null +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -0,0 +1,15 @@ +# 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. + +git clone https://github.com/Project-MONAI/GenerativeModels.git +cd GenerativeModels/ +python setup.py install +cd .. diff --git a/ci/verify_bundle.py b/ci/verify_bundle.py index a337283c..515e339e 100644 --- a/ci/verify_bundle.py +++ b/ci/verify_bundle.py @@ -11,6 +11,7 @@ import argparse import os +import subprocess import sys from typing import List @@ -19,6 +20,7 @@ exclude_verify_preferred_files_list, exclude_verify_shape_list, exclude_verify_torchscript_list, + install_dependency_dict, ) from monai.bundle import ckpt_export, verify_metadata, verify_net_in_out from monai.bundle.config_parser import ConfigParser @@ -292,6 +294,12 @@ def verify_torchscript(bundle_path: str, net_id: str, config_file: str): def verify(bundle, models_path="models", mode="full"): print(f"start verifying {bundle}:") + # install extra dependencies if needed + if bundle in install_dependency_dict.keys(): + script_path = install_dependency_dict[bundle] + install_cmd = f"bash {script_path}" + call_status = subprocess.run(install_cmd, shell=True) + call_status.check_returncode() # add bundle path to ensure custom code can be used sys.path = [os.path.join(models_path, bundle)] + sys.path # verify bundle directory diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index d97266a5..9bc32faf 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -9,7 +9,7 @@ "numpy_version": "1.22.2", "optional_packages_version": { "nibabel": "5.1.0", - "generative": "0.1.0" + "lpips": "0.1.4" }, "task": "BraTS MRI image synthesis", "description": "A generative model for creating 3D brain MRI from Gaussian noise based on BraTS dataset", From c4f8a8a606b9c491326bb90e06ff902d337851c7 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Wed, 26 Apr 2023 14:40:30 +0800 Subject: [PATCH 048/130] use fixed commit id and ignore pytype import error Signed-off-by: Yiheng Wang --- .../install_brats_mri_generative_diffusion_dependency.sh | 1 + pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index 08dd23e8..d9f6bd48 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -11,5 +11,6 @@ git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ +git checkout ad973b8157130eaa146f35628f96eebd2b9d207d python setup.py install cd .. diff --git a/pyproject.toml b/pyproject.toml index ac0df5bf..223ea491 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ check_parameter_types = true # Check variable values against their annotations. check_variable_types = true # Comma or space separated list of error names to ignore. -disable = ["pyi-error"] +disable = ["pyi-error", "import-error"] # Report errors. report_errors = true # Experimental: Infer precise return types even for invalid function calls. From c44dfff0d1835fbc665663f1dd5d4ee72d6c234a Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Wed, 26 Apr 2023 14:13:19 -0700 Subject: [PATCH 049/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 60185215..b2f8e907 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -1,7 +1,7 @@ # Model Overview A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative Model. -This model is trained on BraTS 2018 data (https://luna16.grand-challenge.org/Home/), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). +This model is trained on BraTS 2018 data (https://www.med.upenn.edu/sbia/brats2018.html), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). ![model workflow placeholder, will update](https://miro.medium.com/v2/resize:fit:1256/1*3Jjlrb-TB8hPpIexoCKpZQ.png) From a5777892093392649ac57cdfd7daa1a2068c3f73 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 27 Apr 2023 10:44:03 -0700 Subject: [PATCH 050/130] let lr to be easier configured Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 3 ++- .../configs/train_diffusion.json | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 350431b7..81467d81 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -11,6 +11,7 @@ "dataset_dir": "$@bundle_root", "download_brats": false, "train_batch_size": 2, + "lr": 1e-05, "train_patch_size": [ 112, 128, @@ -69,7 +70,7 @@ "doptimizer": { "_target_": "torch.optim.Adam", "params": "$@dnetwork.parameters()", - "lr": 1e-05 + "lr": "@lr" }, "goptimizer": { "_target_": "torch.optim.Adam", diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 32fae5a7..42bf0926 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -1,6 +1,7 @@ { "ckpt_dir": "$@bundle_root + '/models'", "train_batch_size": 4, + "lr": 1e-05, "train_patch_size": [ 144, 176, @@ -41,7 +42,7 @@ "optimizer": { "_target_": "torch.optim.Adam", "params": "$@diffusion.parameters()", - "lr": 1e-05 + "lr": "@lr" }, "lr_scheduler": { "_target_": "torch.optim.lr_scheduler.MultiStepLR", From ab686887c31d3f38b94695ce0f40226d7a2756cc Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 27 Apr 2023 10:45:07 -0700 Subject: [PATCH 051/130] let lr to be easier configured Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 81467d81..cd9e49cf 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -75,7 +75,7 @@ "goptimizer": { "_target_": "torch.optim.Adam", "params": "$@gnetwork.parameters()", - "lr": 1e-05 + "lr": "@lr" }, "preprocessing_transforms": [ { From f63f326d19eeaf72607308467b736b83757eb37d Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 27 Apr 2023 17:53:24 -0700 Subject: [PATCH 052/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index b2f8e907..c18363aa 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -88,9 +88,9 @@ If it is already downloaded, check if `"dataset_dir"` in `configs/train_autoenco python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -Or run it with multi-gpu: +Or run it with multi-gpu, which requires the learning rate to be scaled up with the number of GPUs. ``` -torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 ``` It take 9 hours when training with 8 32G GPU. @@ -100,9 +100,9 @@ python -m monai.bundle run --config_file "['configs/train_autoencoder.json','con ``` It will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. -Or run it with multi-gpu: +Or run it with multi-gpu, which requires the learning rate to be scaled up with the number of GPUs. ``` -torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` From 364af38183fba94062f15cffdfb0f26aa390ae04 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 27 Apr 2023 17:56:58 -0700 Subject: [PATCH 053/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index c18363aa..1b36a922 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -88,7 +88,7 @@ If it is already downloaded, check if `"dataset_dir"` in `configs/train_autoenco python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -Or run it with multi-gpu, which requires the learning rate to be scaled up with the number of GPUs. +Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 ``` @@ -100,7 +100,7 @@ python -m monai.bundle run --config_file "['configs/train_autoencoder.json','con ``` It will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. -Or run it with multi-gpu, which requires the learning rate to be scaled up with the number of GPUs. +Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` From 91cb9159e5bde992c7cd10ce3021081435eb5e5e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Apr 2023 00:57:08 +0000 Subject: [PATCH 054/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 1b36a922..929ee0e7 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -88,7 +88,7 @@ If it is already downloaded, check if `"dataset_dir"` in `configs/train_autoenco python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. +Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 ``` From 081dcd405d4845d1834c1d506f3141c3af9b6a49 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 30 Apr 2023 22:21:50 -0700 Subject: [PATCH 055/130] add trained weights Signed-off-by: Can-Zhao --- .../docs/README.md | 20 ++----------------- .../large_files.yml | 13 ++++++++++++ 2 files changed, 15 insertions(+), 18 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/large_files.yml diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 929ee0e7..dae5fd47 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -115,31 +115,15 @@ The generated image will be saved to `./output/0` ### Export -The autoencoder and latent diffusion generators can be exported to a Torchscript bundle with the following: - +The autoencoder can be exported to a Torchscript bundle with the following: ``` python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" -python -m monai.bundle ckpt_export diffusion_def --filepath models/model_ldm.ts --ckpt_file models/model_ldm.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` -The models can be loaded after this operation. For example, an image can be generated from a set of random values with: - +The models can be loaded after this operation: ```python import torch -from scripts.ldm_sampler import LDMSampler -from generative.networks.schedulers import DDIMScheduler autoencoder = torch.jit.load("models/model_autoencoder.ts") -diffusion = torch.jit.load("models/model_ldm.ts") -noise = torch.rand(1, 8,36,44,28).to(device) -noise_scheduler = DDIMScheduler( - num_train_timesteps=1000, - beta_start= 0.0015, - beta_end=0.0195, - beta_schedule= "scaled_linear", - clip_sample= False) -noise_scheduler.set_timesteps(num_inference_steps=50) -sampler = LDMSampler() -img = sampler.sampling_fn(noise, autoencoder, diffusion, noise_scheduler) ``` # License diff --git a/models/brats_mri_generative_diffusion/large_files.yml b/models/brats_mri_generative_diffusion/large_files.yml new file mode 100644 index 00000000..4a0a5ab0 --- /dev/null +++ b/models/brats_mri_generative_diffusion/large_files.yml @@ -0,0 +1,13 @@ +large_files: + - path: "models/model_autoencoder.pt" + url: "https://drive.google.com/uc?id=1arp3w8glsQw2h7mQBbk71krqmaG_std6" + hash_val: "314c9091a4b101461a5cc305995a4565" + hash_type: "md5" + - path: "models/model_ldm.pt" + url: "https://drive.google.com/uc?id=1m2pcbj8NMoxEIAOmD9dgYBN4gNcrMx6e" + hash_val: "2463887ef1b704494014e2aecb72802b" + hash_type: "md5" + - path: "models/model_autoencoder.ts" + url: "https://drive.google.com/uc?id=17WybvJHfxL-jc5toFFHxCfgVmFun-5ul" + hash_val: "8ad34a1080a71cf0a21964621a4f3e2e" + hash_type: "md5" From 75273d25d6cda22901408f829b780162ec0096b7 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 30 Apr 2023 22:24:16 -0700 Subject: [PATCH 056/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index dae5fd47..d306eb9c 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -12,15 +12,18 @@ This is a demonstration network meant to just show the training process for this ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` +pip install lpips +pip install git+https://github.com/Project-MONAI/GenerativeModels.git#egg=Generative +``` + +Or we can install it from source +``` +pip install lpips git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ python setup.py install cd .. ``` -We also need -``` -pip install lpips -``` ## Downloading the Dataset The training data is from the [Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018](https://www.med.upenn.edu/sbia/brats2018.html). From b25ab6313224e954463f94eb4036359e745dd2f8 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 30 Apr 2023 22:56:40 -0700 Subject: [PATCH 057/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d306eb9c..c8b41219 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative This model is trained on BraTS 2018 data (https://www.med.upenn.edu/sbia/brats2018.html), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). -![model workflow placeholder, will update](https://miro.medium.com/v2/resize:fit:1256/1*3Jjlrb-TB8hPpIexoCKpZQ.png) +![model workflow](https://miro.medium.com/v2/resize:fit:1256/1*3Jjlrb-TB8hPpIexoCKpZQ.png) This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. @@ -108,6 +108,11 @@ Or run it with multi-gpu, which requires the learning rate to be scaled up accor torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` +

+ autoencoder training curve +        + latent diffusion training curve +

### Inference ``` From 70daa8899b72953580135ac50f08bd7bf7d56b14 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 30 Apr 2023 23:04:27 -0700 Subject: [PATCH 058/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index c8b41219..18e8c9ab 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -115,11 +115,16 @@ torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config

### Inference +The following code generates a synthetic image from a random sampled noise. ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` The generated image will be saved to `./output/0` +An example putput is shown below. Note that this is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. + +![Example synthetic image](placeholder) + ### Export From 20900dc026f683173e6cc219a76358fb1c45f413 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 30 Apr 2023 23:04:52 -0700 Subject: [PATCH 059/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 18e8c9ab..c6d47851 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -121,7 +121,7 @@ python -m monai.bundle run --config_file "['configs/train_autoencoder.json','con ``` The generated image will be saved to `./output/0` -An example putput is shown below. Note that this is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. +An example output is shown below. Note that this is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. ![Example synthetic image](placeholder) From bf4a89bca95b8ebfc8b69c935d5a0437f46f124a Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 1 May 2023 13:54:14 -0700 Subject: [PATCH 060/130] update readme figs Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index c6d47851..b6e47322 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -3,7 +3,7 @@ A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative This model is trained on BraTS 2018 data (https://www.med.upenn.edu/sbia/brats2018.html), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). -![model workflow](https://miro.medium.com/v2/resize:fit:1256/1*3Jjlrb-TB8hPpIexoCKpZQ.png) +![model workflow](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_network.png) This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. @@ -109,9 +109,9 @@ torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config ```

- autoencoder training curve + autoencoder training curve         - latent diffusion training curve + latent diffusion training curve

### Inference @@ -123,7 +123,7 @@ The generated image will be saved to `./output/0` An example output is shown below. Note that this is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. -![Example synthetic image](placeholder) +![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) ### Export From 118c423cece8a53b7ae1ab2f3f7681763f69f465 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 1 May 2023 15:38:39 -0700 Subject: [PATCH 061/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/inference.json | 2 +- .../configs/train_diffusion.json | 2 +- models/brats_mri_generative_diffusion/docs/README.md | 4 +--- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 8463c157..82a1d212 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -10,7 +10,7 @@ "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", - "load_diffusion_path": "$@model_dir + '/model_ldm.pt'", + "load_diffusion_path": "$@model_dir + '/model.pt'", "load_diffusion": "$@diffusion_def.load_state_dict(torch.load(@load_diffusion_path))", "diffusion": "$@diffusion_def.to(@device)", "noise_scheduler": { diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 42bf0926..7d6863d8 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -116,7 +116,7 @@ "save_interval": 0, "save_final": true, "epoch_level": true, - "final_filename": "model_ldm.pt" + "final_filename": "model.pt" }, { "_target_": "StatsHandler", diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index b6e47322..830167f4 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2018 data (https://www.med.upenn.edu/sbia/brats20 This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. +**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers.** ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by @@ -121,8 +121,6 @@ python -m monai.bundle run --config_file "['configs/train_autoencoder.json','con ``` The generated image will be saved to `./output/0` -An example output is shown below. Note that this is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers. - ![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) From ecec425eb7f8d09f0292c63f7b9c1a41e8572d4f Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 1 May 2023 15:47:55 -0700 Subject: [PATCH 062/130] change file name to meet test Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/large_files.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/large_files.yml b/models/brats_mri_generative_diffusion/large_files.yml index 4a0a5ab0..24af4605 100644 --- a/models/brats_mri_generative_diffusion/large_files.yml +++ b/models/brats_mri_generative_diffusion/large_files.yml @@ -3,7 +3,7 @@ large_files: url: "https://drive.google.com/uc?id=1arp3w8glsQw2h7mQBbk71krqmaG_std6" hash_val: "314c9091a4b101461a5cc305995a4565" hash_type: "md5" - - path: "models/model_ldm.pt" + - path: "models/model.pt" url: "https://drive.google.com/uc?id=1m2pcbj8NMoxEIAOmD9dgYBN4gNcrMx6e" hash_val: "2463887ef1b704494014e2aecb72802b" hash_type: "md5" From de4abf44c843a46ddafa509d62fa34a781f68ae6 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 1 May 2023 16:05:54 -0700 Subject: [PATCH 063/130] update meta Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/metadata.json | 1 + 1 file changed, 1 insertion(+) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 9bc32faf..2bbf89db 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -11,6 +11,7 @@ "nibabel": "5.1.0", "lpips": "0.1.4" }, + "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", "authors": "MONAI team", From 28896e3cba1cca66efb5235e7cfdf67a5df91ae9 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 2 May 2023 01:15:01 -0700 Subject: [PATCH 064/130] update dataset in readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 830167f4..20d60f03 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -1,11 +1,11 @@ # Model Overview A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative Model. -This model is trained on BraTS 2018 data (https://www.med.upenn.edu/sbia/brats2018.html), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). +This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http://medicaldecathlon.com/), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). ![model workflow](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_network.png) -This model is a generator for creating images like the Flair MRIs based on BraTS 2018 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. +This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. **This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers.** @@ -26,7 +26,7 @@ cd .. ``` ## Downloading the Dataset -The training data is from the [Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018](https://www.med.upenn.edu/sbia/brats2018.html). +The training data is Brats 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). Target: image generatiion Task: Synthesis From c6c7416a206e214913c059ed783e0ac03e3e6b5b Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 4 May 2023 12:31:03 -0700 Subject: [PATCH 065/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 20d60f03..3d190769 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -31,7 +31,7 @@ The training data is Brats 2016 and 2017 from [Medical Decathlon](http://medical Target: image generatiion Task: Synthesis Modality: MRI -Size: 285 3D volumes (1 channel used) +Size: 388 3D volumes (1 channel used) The dataset can be downloaded automatically at the beggining of training. @@ -49,7 +49,7 @@ The training of autoencoder was performed with the following: - Loss: L1 loss, perceptual loss, KL divergence loss, adversianl loss, GAN BCE loss #### Input -1 channel 3D MRI patches +1 channel 3D MRI Flair patches #### Output - 1 channel 3D MRI reconstructed patches From f2ed6879c93a9ddd75bf6064321834a828638c4f Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 4 May 2023 12:31:33 -0700 Subject: [PATCH 066/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 3d190769..2aa568fa 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -95,7 +95,7 @@ Or run it with multi-gpu, which requires the learning rate to be scaled up accor ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 ``` -It take 9 hours when training with 8 32G GPU. +It take 9 hours when training with 8 GPU each using 32G memory. After the autoencoder is trained, run the following command to train the latent diffusion model. ``` From fd1ee286c56c29c4fcc5c30826b44ce673dce6fb Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 4 May 2023 12:35:36 -0700 Subject: [PATCH 067/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 2aa568fa..d76dda2f 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -13,7 +13,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` pip install lpips -pip install git+https://github.com/Project-MONAI/GenerativeModels.git#egg=Generative +pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 ``` Or we can install it from source From ff0a45333a97c05b24a4696903587b19a2e1d09e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 20:38:56 -0700 Subject: [PATCH 068/130] change diffusion_def to network_def Signed-off-by: Can-Zhao --- .../configs/train_diffusion.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index 7d6863d8..b3c76147 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -16,7 +16,7 @@ "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", "autoencoder": "$@autoencoder_def.to(@device)", - "diffusion_def": { + "network_def": { "_target_": "generative.networks.nets.DiffusionModelUNet", "spatial_dims": "@spatial_dims", "in_channels": "@latent_channels", @@ -38,7 +38,7 @@ ], "num_res_blocks": 2 }, - "diffusion": "$@diffusion_def.to(@device)", + "diffusion": "$@network_def.to(@device)", "optimizer": { "_target_": "torch.optim.Adam", "params": "$@diffusion.parameters()", From 5d5d62b1f938461804312a357a416fc35f3d07d4 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 20:40:43 -0700 Subject: [PATCH 069/130] change sampler to inferer Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/inference.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 82a1d212..8629d4cd 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -31,14 +31,14 @@ "_target_": "scripts.ldm_sampler.LDMSampler", "_requires_": "@set_timesteps" }, - "sample": "$@sampler.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", + "inferer": "$@sampler.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", "saver": { "_target_": "SaveImage", "_requires_": "@create_output_dir", "output_dir": "@output_dir", "output_postfix": "@output_postfix" }, - "generated_image": "$@sample", + "generated_image": "$@inferer", "run": [ "$@saver(@generated_image[0])" ] From 71666399f854ed7f12a5bd9c25a4a7ca0a6fff2c Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 20:41:22 -0700 Subject: [PATCH 070/130] change sampler to inferer Signed-off-by: Can-Zhao --- .../brats_mri_generative_diffusion/configs/inference.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 8629d4cd..3d2daa08 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -27,18 +27,18 @@ }, "noise": "$torch.randn([1]+@latent_shape).to(@device)", "set_timesteps": "$@noise_scheduler.set_timesteps(num_inference_steps=50)", - "sampler": { + "inferer": { "_target_": "scripts.ldm_sampler.LDMSampler", "_requires_": "@set_timesteps" }, - "inferer": "$@sampler.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", + "sample": "$@inferer.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", "saver": { "_target_": "SaveImage", "_requires_": "@create_output_dir", "output_dir": "@output_dir", "output_postfix": "@output_postfix" }, - "generated_image": "$@inferer", + "generated_image": "$@sample", "run": [ "$@saver(@generated_image[0])" ] From e78eb49aa5aad123f0a2cf305b6a82a23287e555 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:01:45 -0700 Subject: [PATCH 071/130] change network_def Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/inference.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 3d2daa08..80b73ea0 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -11,8 +11,8 @@ "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", "load_diffusion_path": "$@model_dir + '/model.pt'", - "load_diffusion": "$@diffusion_def.load_state_dict(torch.load(@load_diffusion_path))", - "diffusion": "$@diffusion_def.to(@device)", + "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))", + "diffusion": "$@network_def.to(@device)", "noise_scheduler": { "_target_": "generative.networks.schedulers.DDIMScheduler", "_requires_": [ From 65b36987a4683dc7c93b5017b2ec13113f1a7a2b Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:03:06 -0700 Subject: [PATCH 072/130] update latent channel Signed-off-by: Can-Zhao --- .../brats_mri_generative_diffusion/configs/train_diffusion.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index b3c76147..b1be476b 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -8,7 +8,7 @@ 112 ], "latent_shape": [ - 8, + "@latent_channels", 36, 44, 28 From 4fd082c1de2eeb5c58cd5b56efd9062eb032c90b Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:10:13 -0700 Subject: [PATCH 073/130] change network_def Signed-off-by: Can-Zhao --- .../configs/multi_gpu_train_diffusion.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json index 89a381a7..fbf130d2 100644 --- a/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/multi_gpu_train_diffusion.json @@ -1,7 +1,7 @@ { "diffusion": { "_target_": "torch.nn.parallel.DistributedDataParallel", - "module": "$@diffusion_def.to(@device)", + "module": "$@network_def.to(@device)", "device_ids": [ "@device" ], From 7d9859539784ed8909e65315ed679e21b2401ae3 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:13:20 -0700 Subject: [PATCH 074/130] change network_def Signed-off-by: Can-Zhao --- .../configs/inference.json | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index 80b73ea0..ffec978a 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -10,6 +10,28 @@ "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", + "network_def": { + "_target_": "generative.networks.nets.DiffusionModelUNet", + "spatial_dims": "@spatial_dims", + "in_channels": "@latent_channels", + "out_channels": "@latent_channels", + "num_channels": [ + 256, + 256, + 512 + ], + "attention_levels": [ + false, + true, + true + ], + "num_head_channels": [ + 0, + 64, + 64 + ], + "num_res_blocks": 2 + }, "load_diffusion_path": "$@model_dir + '/model.pt'", "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))", "diffusion": "$@network_def.to(@device)", From 5340af429351f01c6ba098ebaf8e6fc6f237f5d5 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:20:18 -0700 Subject: [PATCH 075/130] add paramaters in inference Signed-off-by: Can-Zhao --- .../configs/inference.json | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index ffec978a..c5a324a1 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -10,6 +10,31 @@ "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", + "spatial_dims": 3, + "image_channels": 1, + "latent_channels": 8, + "autoencoder_def": { + "_target_": "generative.networks.nets.AutoencoderKL", + "spatial_dims": "@spatial_dims", + "in_channels": "@image_channels", + "out_channels": "@image_channels", + "latent_channels": "@latent_channels", + "num_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 + }, "network_def": { "_target_": "generative.networks.nets.DiffusionModelUNet", "spatial_dims": "@spatial_dims", @@ -32,6 +57,9 @@ ], "num_res_blocks": 2 }, + "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", + "load_autoencoder": "$@autoencoder_def.load_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))", "diffusion": "$@network_def.to(@device)", From 8342dde1bcf1e554a3801f384c2f3719b0d84494 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 5 May 2023 21:31:41 -0700 Subject: [PATCH 076/130] update meta Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/metadata.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 2bbf89db..f9dc18c3 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -1,5 +1,5 @@ { - "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json", + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_20220718.json", "version": "1.0.0", "changelog": { "1.0.0": "Initial release" From 2602221cb8c720efbacf0bc11026aa9efdaa6cd8 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sat, 6 May 2023 21:54:34 -0700 Subject: [PATCH 077/130] update meta Signed-off-by: Can-Zhao --- .../configs/metadata.json | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index f9dc18c3..b5f9c112 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -28,6 +28,11 @@ "type": "image", "format": "random noise", "num_channels": 1, + "spatial_shape": [ + 36, + 44, + 28 + ], "dtype": "float32", "value_range": [ 0, @@ -39,13 +44,13 @@ "outputs": { "pred": { "type": "image", - "format": "magnitude", - "modality": "MR", + "format": "feature map", + "modality": "feature map of MR", "num_channels": 1, "spatial_shape": [ - 144, - 176, - 112 + 36, + 44, + 28 ], "dtype": "float32", "value_range": [ From af92fb184a7a2c1285f147709a2db3460e713e29 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sat, 6 May 2023 22:38:12 -0700 Subject: [PATCH 078/130] update readme Signed-off-by: Can-Zhao --- .../configs/metadata.json | 17 ++++++++--------- .../docs/README.md | 4 ++-- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index b5f9c112..3c3e706e 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -24,10 +24,10 @@ "references": [], "network_data_format": { "inputs": { - "image": { + "latent": { "type": "image", - "format": "random noise", - "num_channels": 1, + "format": "image", + "num_channels": 8, "spatial_shape": [ 36, 44, @@ -38,15 +38,14 @@ 0, 1 ], - "is_patch_data": false + "is_patch_data": true } }, "outputs": { "pred": { "type": "image", - "format": "feature map", - "modality": "feature map of MR", - "num_channels": 1, + "format": "image", + "num_channels": 8, "spatial_shape": [ 36, 44, @@ -57,9 +56,9 @@ 0, 1 ], - "is_patch_data": false, + "is_patch_data": true, "channel_def": { - "0": "Flair" + "0": "image" } } } diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d76dda2f..edf78cec 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to have GPU with memory larger than 32G to enable larger networks and attention layers.** +**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like Brats21 and have GPU with memory larger than 32G to enable larger networks and attention layers.** ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by @@ -117,7 +117,7 @@ torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config ### Inference The following code generates a synthetic image from a random sampled noise. ``` -python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" +python -m monai.bundle run --config_file configs/inference.json ``` The generated image will be saved to `./output/0` From 2cbf81aeba6d31f50baadd9d34acaf6d5a065f6e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 7 May 2023 20:57:08 -0700 Subject: [PATCH 079/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index edf78cec..e961d78d 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like Brats21 and have GPU with memory larger than 32G to enable larger networks and attention layers.** +**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by From e17105aabc44faacafcdf0beb616649692826914 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 7 May 2023 22:25:12 -0700 Subject: [PATCH 080/130] update meta Signed-off-by: Can-Zhao --- .../configs/metadata.json | 49 +++++++++++++++++-- 1 file changed, 45 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 3c3e706e..c95f8923 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -1,5 +1,5 @@ { - "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_20220718.json", + "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20230507.json", "version": "1.0.0", "changelog": { "1.0.0": "Initial release" @@ -22,11 +22,52 @@ "eval_metrics": {}, "intended_use": "This is a research tool/prototype and not to be used clinically", "references": [], - "network_data_format": { + "autoencoder_data_format": { "inputs": { - "latent": { + "image": { "type": "image", "format": "image", + "num_channels": 1, + "spatial_shape": [ + 112, + 128, + 80 + ], + "dtype": "float32", + "value_range": [ + 0, + 1 + ], + "is_patch_data": true + } + }, + "outputs": { + "pred": { + "type": "image", + "format": "image", + "num_channels": 1, + "spatial_shape": [ + 112, + 128, + 80 + ], + "dtype": "float32", + "value_range": [ + 0, + 1 + ], + "is_patch_data": true, + "channel_def": { + "0": "image" + } + } + } + }, + "generator_data_format": { + "inputs": { + "latent": { + "type": "noise", + "format": "image", "num_channels": 8, "spatial_shape": [ 36, @@ -43,7 +84,7 @@ }, "outputs": { "pred": { - "type": "image", + "type": "feature", "format": "image", "num_channels": 8, "spatial_shape": [ From d4b6020b471c666bb2d9fdd0b7b53bc41e414a69 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 7 May 2023 22:40:30 -0700 Subject: [PATCH 081/130] update meta Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/metadata.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index c95f8923..bbbc3b16 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -5,7 +5,7 @@ "1.0.0": "Initial release" }, "monai_version": "1.2.0rc5", - "pytorch_version": "2.0.0", + "pytorch_version": "1.13.1", "numpy_version": "1.22.2", "optional_packages_version": { "nibabel": "5.1.0", @@ -16,7 +16,7 @@ "description": "A generative model for creating 3D brain MRI from Gaussian noise based on BraTS dataset", "authors": "MONAI team", "copyright": "Copyright (c) MONAI Consortium", - "data_source": "https://www.med.upenn.edu/sbia/brats2018/data.html", + "data_source": "http://medicaldecathlon.com/", "data_type": "nibabel", "image_classes": "Flair brain MRI with 1.1x1.1x1.1 mm voxel size", "eval_metrics": {}, From 54bf9f07bb717e97e4eef0c10f466107b81250bb Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 7 May 2023 22:43:02 -0700 Subject: [PATCH 082/130] update inference.json Signed-off-by: Can-Zhao --- .../brats_mri_generative_diffusion/configs/inference.json | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/models/brats_mri_generative_diffusion/configs/inference.json b/models/brats_mri_generative_diffusion/configs/inference.json index c5a324a1..b2965eea 100644 --- a/models/brats_mri_generative_diffusion/configs/inference.json +++ b/models/brats_mri_generative_diffusion/configs/inference.json @@ -13,6 +13,12 @@ "spatial_dims": 3, "image_channels": 1, "latent_channels": 8, + "latent_shape": [ + 8, + 36, + 44, + 28 + ], "autoencoder_def": { "_target_": "generative.networks.nets.AutoencoderKL", "spatial_dims": "@spatial_dims", From 6c80f6de2168804c874910c81c193d78d65a688e Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 8 May 2023 15:12:28 +0800 Subject: [PATCH 083/130] skip input check Signed-off-by: Yiheng Wang --- ci/bundle_custom_data.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ci/bundle_custom_data.py b/ci/bundle_custom_data.py index e0485494..dbf91530 100644 --- a/ci/bundle_custom_data.py +++ b/ci/bundle_custom_data.py @@ -13,7 +13,12 @@ # This list is used for our CI tests to determine whether a bundle needs to be tested with # the `verify_data_shape` function in `verify_bundle.py`. # If a bundle does not need to be tested, please add the bundle name into the list. -exclude_verify_shape_list = ["mednist_gan", "lung_nodule_ct_detection", "pathology_nuclei_segmentation_classification"] +exclude_verify_shape_list = [ + "mednist_gan", + "lung_nodule_ct_detection", + "pathology_nuclei_segmentation_classification", + "brats_mri_generative_diffusion", +] # This list is used for our CI tests to determine whether a bundle contains the preferred files. # If a bundle does not have any of the preferred files, please add the bundle name into the list. @@ -28,6 +33,7 @@ "wholeBrainSeg_Large_UNEST_segmentation", "breast_density_classification", "mednist_reg", + "brats_mri_generative_diffusion", ] # This dict is used for our CI tests to install required dependencies that cannot be installed by `pip install` directly. From 8e472f679feb9c1fe927d8c047af0fe71672bddc Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 8 May 2023 00:28:47 -0700 Subject: [PATCH 084/130] update meta Signed-off-by: Can-Zhao --- .../configs/metadata.json | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index bbbc3b16..63380bb9 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -80,6 +80,18 @@ 1 ], "is_patch_data": true + }, + "condition": { + "type": "timesteps", + "format": "timesteps", + "num_channels": 1, + "spatial_shape": [], + "dtype": "float32", + "value_range": [ + 0, + 1000 + ], + "is_patch_data": false } }, "outputs": { From 8a1a142e8bc2e043deae871eb906546108979002 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 8 May 2023 00:30:54 -0700 Subject: [PATCH 085/130] update meta Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/configs/metadata.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/metadata.json b/models/brats_mri_generative_diffusion/configs/metadata.json index 63380bb9..4bcc437d 100644 --- a/models/brats_mri_generative_diffusion/configs/metadata.json +++ b/models/brats_mri_generative_diffusion/configs/metadata.json @@ -86,7 +86,7 @@ "format": "timesteps", "num_channels": 1, "spatial_shape": [], - "dtype": "float32", + "dtype": "long", "value_range": [ 0, 1000 From 969f7e83456f90e6586c4cbc7240f9b41833e97a Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 8 May 2023 07:13:04 -0700 Subject: [PATCH 086/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index e961d78d..af9cec12 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -56,7 +56,7 @@ The training of autoencoder was performed with the following: - 8 channel mean of latent features - 8 channel standard deviation of latent features -### Training configuration of difusion model +### Training configuration of diffusion model The training of latent diffusion model was performed with the following: - GPU: at least 32GB GPU memory @@ -68,6 +68,7 @@ The training of latent diffusion model was performed with the following: #### Training Input 8 channel noisy latent features +a int that indicates the time step #### Training Output 8 channel predicted added noise From a0b69aac6ad6f585cbcff713f9b61880f59c0dd5 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 8 May 2023 07:14:24 -0700 Subject: [PATCH 087/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index af9cec12..459341cf 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -68,7 +68,7 @@ The training of latent diffusion model was performed with the following: #### Training Input 8 channel noisy latent features -a int that indicates the time step +an int that indicates the time step #### Training Output 8 channel predicted added noise From 382be8510d6339914fc663741394ef8ada2d64e9 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Mon, 8 May 2023 07:14:53 -0700 Subject: [PATCH 088/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 459341cf..74d5c199 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -67,8 +67,8 @@ The training of latent diffusion model was performed with the following: - Loss: MSE loss #### Training Input -8 channel noisy latent features -an int that indicates the time step +- 8 channel noisy latent features +- an int that indicates the time step #### Training Output 8 channel predicted added noise From 0f2879d0f718a9d089ed1d12a51f8abee8e657da Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 9 May 2023 16:32:04 +0800 Subject: [PATCH 089/130] reformat readme Signed-off-by: Yiheng Wang --- ...ats_mri_generative_diffusion_dependency.sh | 6 +- .../docs/README.md | 71 +++++++++++-------- 2 files changed, 41 insertions(+), 36 deletions(-) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index d9f6bd48..b3e2a31a 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -9,8 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -git clone https://github.com/Project-MONAI/GenerativeModels.git -cd GenerativeModels/ -git checkout ad973b8157130eaa146f35628f96eebd2b9d207d -python setup.py install -cd .. +pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 74d5c199..defaf940 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -1,7 +1,7 @@ # Model Overview A pre-trained model for volumetric (3D) Brats MRI 3D Latent Diffusion Generative Model. -This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http://medicaldecathlon.com/), using the Latent diffusion model (Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf). +This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http://medicaldecathlon.com/), using the Latent diffusion model [1]. ![model workflow](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_network.png) @@ -79,64 +79,73 @@ The training of latent diffusion model was performed with the following: #### Inference Output 8 channel denoised latent features +### Memory Consumption Warning + +If you face memory issues with data loading, you can lower the caching rate `cache_rate` in the configurations within range [0, 1] to minimize the System RAM requirements. + +## Performance + +#### Training Loss +![A graph showing the autoencoder training curve](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_train_autoencoder_loss.png) + +![A graph showing the latent diffusion training curve](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_train_diffusion_loss.png) + +#### Example synthetic image +![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) ## MONAI Bundle Commands -Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs. -If the Brats dataset is not downloaded, run the following command, the data will be downloaded and extracted to `./Task01_BrainTumour`. -``` -python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True -``` -If it is already downloaded, check if `"dataset_dir"` in `configs/train_autoencoder.json`has folder `Task01_BrainTumour`. If so, run: +In addition to the Pythonic APIs, a few command line interfaces (CLI) are provided to interact with the bundle. The CLI supports flexible use cases, such as overriding configs at runtime and predefining arguments in a file. + +For more details usage instructions, visit the [MONAI Bundle Configuration Page](https://docs.monai.io/en/latest/config_syntax.html). + +#### Execute training: + +- Train autoencoder ``` python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. -``` -torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 -``` -It take 9 hours when training with 8 GPU each using 32G memory. +Please specify "download_brats" into `True` if the dataset is not downloaded. -After the autoencoder is trained, run the following command to train the latent diffusion model. +- Train latent diffusion model ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` + It will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. -Or run it with multi-gpu, which requires the learning rate to be scaled up according to the number of GPUs. +#### Override the `train` config to execute multi-GPU autoencoder training: + +Run with multi-GPU requires the learning rate to be scaled up according to the number of GPUs. + +- Train autoencoder + +``` +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 +``` + +- Train latent difussion model + ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` -

- autoencoder training curve -        - latent diffusion training curve -

+#### Execute inference: -### Inference The following code generates a synthetic image from a random sampled noise. ``` python -m monai.bundle run --config_file configs/inference.json ``` -The generated image will be saved to `./output/0` - -![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) - -### Export +#### Export checkpoint to TorchScript file: -The autoencoder can be exported to a Torchscript bundle with the following: ``` python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` -The models can be loaded after this operation: -```python -import torch -autoencoder = torch.jit.load("models/model_autoencoder.ts") -``` +# References +[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf # License Copyright (c) MONAI Consortium From 38b5e1ee63d35245fb326e7aa89523e5684f6500 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 9 May 2023 17:51:33 +0800 Subject: [PATCH 090/130] enable torchscript test Signed-off-by: Yiheng Wang --- ci/bundle_custom_data.py | 1 - ...ats_mri_generative_diffusion_dependency.sh | 6 +++- ci/verify_bundle.py | 28 +++++++++++++++---- .../docs/README.md | 2 ++ 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/ci/bundle_custom_data.py b/ci/bundle_custom_data.py index dbf91530..20923eab 100644 --- a/ci/bundle_custom_data.py +++ b/ci/bundle_custom_data.py @@ -33,7 +33,6 @@ "wholeBrainSeg_Large_UNEST_segmentation", "breast_density_classification", "mednist_reg", - "brats_mri_generative_diffusion", ] # This dict is used for our CI tests to install required dependencies that cannot be installed by `pip install` directly. diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index b3e2a31a..d9f6bd48 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -9,4 +9,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 +git clone https://github.com/Project-MONAI/GenerativeModels.git +cd GenerativeModels/ +git checkout ad973b8157130eaa146f35628f96eebd2b9d207d +python setup.py install +cd .. diff --git a/ci/verify_bundle.py b/ci/verify_bundle.py index 515e339e..f6b11c2a 100644 --- a/ci/verify_bundle.py +++ b/ci/verify_bundle.py @@ -49,6 +49,20 @@ def _find_bundle_file(root_dir: str, file: str, suffix=("json", "yaml", "yml")): return file_name +def _get_weights_names(bundle: str): + # TODO: this function is temporarily used. It should be replaced by detailed config tests. + if bundle == "brats_mri_generative_diffusion": + return "model_autoencoder.pt", "model_autoencoder.ts" + return "model.pt", "model.ts" + + +def _get_net_id(bundle: str): + # TODO: this function is temporarily used. It should be replaced by detailed config tests. + if bundle == "brats_mri_generative_diffusion": + return "autoencoder_def" + return "network_def" + + def _check_missing_keys(file_name: str, bundle_path: str, keys_list: List): config = ConfigParser.load_config_file(os.path.join(bundle_path, "configs", file_name)) missing_keys = [] @@ -269,7 +283,9 @@ def verify_data_shape(bundle_path: str, net_id: str, config_file: str): ) -def verify_torchscript(bundle_path: str, net_id: str, config_file: str): +def verify_torchscript( + bundle_path: str, net_id: str, config_file: str, model_name: str = "model.pt", ts_name: str = "model.ts" +): """ This function is used to verify if the checkpoint is able to export into torchscript model, and if "models/model.ts" is provided, it will be checked if it is able to be loaded @@ -279,14 +295,14 @@ def verify_torchscript(bundle_path: str, net_id: str, config_file: str): ckpt_export( net_id=net_id, filepath=os.path.join(bundle_path, "models/verify_model.ts"), - ckpt_file=os.path.join(bundle_path, "models/model.pt"), + ckpt_file=os.path.join(bundle_path, "models", model_name), meta_file=os.path.join(bundle_path, "configs/metadata.json"), config_file=os.path.join(bundle_path, config_file), bundle_root=bundle_path, ) print("export weights into TorchScript module successfully.") - ts_model_path = os.path.join(bundle_path, "models/model.ts") + ts_model_path = os.path.join(bundle_path, "models", ts_name) if os.path.exists(ts_model_path): torch.jit.load(ts_model_path) print("Provided TorchScript module is verified correctly.") @@ -321,7 +337,8 @@ def verify(bundle, models_path="models", mode="full"): return # The following are optional tests and require GPU - net_id, inference_file_name = "network_def", _find_bundle_file(os.path.join(bundle_path, "configs"), "inference") + net_id = _get_net_id(bundle) + inference_file_name = _find_bundle_file(os.path.join(bundle_path, "configs"), "inference") config_file = os.path.join("configs", inference_file_name) if bundle in exclude_verify_shape_list: @@ -333,7 +350,8 @@ def verify(bundle, models_path="models", mode="full"): if bundle in exclude_verify_torchscript_list: print(f"bundle: {bundle} does not support torchscript, skip verifying.") else: - verify_torchscript(bundle_path, net_id, config_file) + model_name, ts_name = _get_weights_names(bundle=bundle) + verify_torchscript(bundle_path, net_id, config_file, model_name, ts_name) if __name__ == "__main__": diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index defaf940..4889e082 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -140,6 +140,8 @@ python -m monai.bundle run --config_file configs/inference.json #### Export checkpoint to TorchScript file: +The autoencoder can be exported into a TorchScript file. + ``` python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" ``` From 66189586d39f4a59e582687ed57677cf9c437aa7 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 9 May 2023 11:37:19 -0700 Subject: [PATCH 091/130] simplify export Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 4889e082..acb5d323 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -143,7 +143,7 @@ python -m monai.bundle run --config_file configs/inference.json The autoencoder can be exported into a TorchScript file. ``` -python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/inference.json']" +python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json ``` # References From c3e907ef5a84f65d043b4aecb20a08baebfd861e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 9 May 2023 11:38:25 -0700 Subject: [PATCH 092/130] add lpips version Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index acb5d323..91349ab8 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -12,7 +12,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS ## Install the dependency of MONAI generative models [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` -pip install lpips +pip install lpips==0.1.4 pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 ``` From 626a4876ac84a69f54ca5f403050071e9d7af037 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 9 May 2023 11:39:15 -0700 Subject: [PATCH 093/130] add lpips version Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 91349ab8..be893094 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -18,7 +18,7 @@ pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 Or we can install it from source ``` -pip install lpips +pip install lpips==0.1.4 git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ python setup.py install From fe96105a73430af48d2987dd1f48076a7e333f8f Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 9 May 2023 16:10:32 -0700 Subject: [PATCH 094/130] update readme Signed-off-by: Can-Zhao --- .../docs/README.md | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index be893094..d636e37b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -9,7 +9,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS **This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** -## Install the dependency of MONAI generative models +## MONAI Generative Model Dependencies [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` pip install lpips==0.1.4 @@ -25,28 +25,28 @@ python setup.py install cd .. ``` -## Downloading the Dataset -The training data is Brats 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). +## Data +The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). -Target: image generatiion -Task: Synthesis -Modality: MRI -Size: 388 3D volumes (1 channel used) +- Target: Image Generation +- Task: Synthesis +- Modality: MRI +- Size: 388 3D volumes (1 channel used) -The dataset can be downloaded automatically at the beggining of training. +The dataset can be downloaded automatically at the beginning of training. ## Training configuration -If user has GPU memory smaller than 32G, then please decrease the `"train_batch_size"` in `configs/train_autoencoder.json` and `configs/train_diffusion.json`. +If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. -### Training configuration of autoencoder -The training of autoencoder was performed with the following: +### Training Configuration of Autoencoder +The autoencoder was trained using the following configuration: - GPU: at least 32GB GPU memory - Actual Model Input: 112 x 128 x 80 - AMP: False - Optimizer: Adam - Learning Rate: 1e-5 -- Loss: L1 loss, perceptual loss, KL divergence loss, adversianl loss, GAN BCE loss +- Loss: L1 loss, perceptual loss, KL divergence loss, adversarial loss, GAN BCE loss #### Input 1 channel 3D MRI Flair patches @@ -56,8 +56,8 @@ The training of autoencoder was performed with the following: - 8 channel mean of latent features - 8 channel standard deviation of latent features -### Training configuration of diffusion model -The training of latent diffusion model was performed with the following: +### Training Configuration of Diffusion Model +The latent diffusion model was trained using the following configuration: - GPU: at least 32GB GPU memory - Actual Model Input: 36 x 44 x 28 From 9779ae79a4515161223d1405d68b3ea80e63520f Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Tue, 9 May 2023 16:11:37 -0700 Subject: [PATCH 095/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d636e37b..d2228148 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -35,7 +35,7 @@ The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medical The dataset can be downloaded automatically at the beginning of training. -## Training configuration +## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. ### Training Configuration of Autoencoder From 1cdfb5fea0039320c628a1b601e533b152daa5d6 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Wed, 10 May 2023 16:31:25 +0800 Subject: [PATCH 096/130] update command format Signed-off-by: Yiheng Wang --- .../docs/README.md | 39 ++++++++++--------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d2228148..4b2ea897 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -9,6 +9,9 @@ This model is a generator for creating images like the Flair MRIs based on BraTS **This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** +#### Example synthetic image +![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) + ## MONAI Generative Model Dependencies [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` @@ -90,49 +93,49 @@ If you face memory issues with data loading, you can lower the caching rate `cac ![A graph showing the latent diffusion training curve](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_train_diffusion_loss.png) -#### Example synthetic image -![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) - ## MONAI Bundle Commands In addition to the Pythonic APIs, a few command line interfaces (CLI) are provided to interact with the bundle. The CLI supports flexible use cases, such as overriding configs at runtime and predefining arguments in a file. For more details usage instructions, visit the [MONAI Bundle Configuration Page](https://docs.monai.io/en/latest/config_syntax.html). -#### Execute training: +#### Execute Autoencoder Training (w/data download) + +Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs using the BraTS dataset. If the dataset is not downloaded, it will be automatically downloaded and extracted to `./Task01_BrainTumour`. -- Train autoencoder ``` -python -m monai.bundle run --config_file configs/train_autoencoder.json +python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True ``` -Please specify "download_brats" into `True` if the dataset is not downloaded. +#### Execute Autoencoder Training +If the dataset is already downloaded, make sure that `"dataset_dir"` in `configs/train_autoencoder.json` has the correct path to the dataset `Task01_BrainTumour`. Then, run: -- Train latent diffusion model ``` -python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" +python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -It will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. - -#### Override the `train` config to execute multi-GPU autoencoder training: +#### Override the `train` config to execute multi-GPU training for Autoencoder: +To train with multiple GPUs, use the following command, which requires scaling up the learning rate according to the number of GPUs. -Run with multi-GPU requires the learning rate to be scaled up according to the number of GPUs. +``` +torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 +``` -- Train autoencoder +#### Execute Latent Diffusion Model Training +After training the autoencoder, run the following command to train the latent diffusion model. This command will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. ``` -torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 +python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` -- Train latent difussion model +#### Override the `train` config to execute multi-GPU training for Latent Diffusion Model: +To train with multiple GPUs, use the following command, which requires scaling up the learning rate according to the number of GPUs. ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` #### Execute inference: - The following code generates a synthetic image from a random sampled noise. ``` python -m monai.bundle run --config_file configs/inference.json @@ -140,7 +143,7 @@ python -m monai.bundle run --config_file configs/inference.json #### Export checkpoint to TorchScript file: -The autoencoder can be exported into a TorchScript file. +The Autoencoder can be exported into a TorchScript file. ``` python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json From d0739b36d4514901ad5473766360014d74b596ba Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 12:27:22 -0700 Subject: [PATCH 097/130] perceptual loss and config Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 9 ++++++--- .../scripts/losses.py | 15 ++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index cd9e49cf..1f69d021 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -8,7 +8,8 @@ "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", - "dataset_dir": "$@bundle_root", + "dataset_dir": "@bundle_root", + "perceptual_loss_pretrained_path": "$@bundle_root +'/nv_resnet50.pth.tar'", "download_brats": false, "train_batch_size": 2, "lr": 1e-05, @@ -60,9 +61,11 @@ "perceptual_loss_def": { "_target_": "generative.losses.PerceptualLoss", "spatial_dims": "@spatial_dims", - "network_type": "squeeze", + "network_type": "resnet50", "is_fake_3d": true, - "fake_3d_ratio": 0.2 + "fake_3d_ratio": 0.2, + "pretrained_path": "@perceptual_loss_pretrained_path", + "pretrained_state_dict_key": "state_dict" }, "dnetwork": "$@discriminator_def.to(@device)", "gnetwork": "$@autoencoder_def.to(@device)", diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 8a0deec5..e2b1343a 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -9,28 +9,25 @@ # See the License for the specific language governing permissions and import torch -from generative.losses import PatchAdversarialLoss, PerceptualLoss +from generative.losses import PatchAdversarialLoss intensity_loss = torch.nn.L1Loss() adv_loss = PatchAdversarialLoss(criterion="least_squares") -loss_perceptual = PerceptualLoss(spatial_dims=3, network_type="squeeze", is_fake_3d=True, fake_3d_ratio=0.2) -adv_weight = 0.01 -perceptual_weight = 0.01 +adv_weight = 0.1 +perceptual_weight = 0.1 # kl_weight: important hyper-parameter. # If too large, decoder cannot recon good results from latent space. # If too small, latent space will not be regularized enough for the diffusion model kl_weight = 1e-7 - -def compute_kl_loss(z_mu, z_sigma): +def KL_loss(z_mu, z_sigma): kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) return torch.sum(kl_loss) / kl_loss.shape[0] - -def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual=loss_perceptual): +def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) - kl_loss = compute_kl_loss(z_mu, z_sigma) + kl_loss = KL_loss(z_mu, z_sigma) p_loss = loss_perceptual(gen_images.float(), real_images.float()) loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss From 0c3e9c99a3a05922e855bc556f43f5f26a7b9be3 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 12:34:09 -0700 Subject: [PATCH 098/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 4b2ea897..51afa724 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -10,6 +10,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS **This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** #### Example synthetic image +An example result from inference is shown below: ![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) ## MONAI Generative Model Dependencies From daf6b674d1a08ad2dee1795fcb24bda0ec899a7e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 12:38:01 -0700 Subject: [PATCH 099/130] reformat Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/losses.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index e2b1343a..199e13d8 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -21,10 +21,12 @@ # If too small, latent space will not be regularized enough for the diffusion model kl_weight = 1e-7 + def KL_loss(z_mu, z_sigma): kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) return torch.sum(kl_loss) / kl_loss.shape[0] + def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) kl_loss = KL_loss(z_mu, z_sigma) From c9b7c821f1ae89d0c0e988e1f622f2be4a784063 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 12:38:52 -0700 Subject: [PATCH 100/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 51afa724..f5e3b365 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,12 +7,12 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** - #### Example synthetic image An example result from inference is shown below: ![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) +**This is a demonstration network meant to just show the training process for this sort of network with MONAI. To achieve better performance, users need to use larger dataset like [Brats 2021](https://www.synapse.org/#!Synapse:syn25829067/wiki/610865) and have GPU with memory larger than 32G to enable larger networks and attention layers.** + ## MONAI Generative Model Dependencies [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` From 7c1e499b9f9e179fa64b8ca5bc63e876c1e36e0e Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 12:52:39 -0700 Subject: [PATCH 101/130] flake8 Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/scripts/losses.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/scripts/losses.py b/models/brats_mri_generative_diffusion/scripts/losses.py index 199e13d8..43536067 100644 --- a/models/brats_mri_generative_diffusion/scripts/losses.py +++ b/models/brats_mri_generative_diffusion/scripts/losses.py @@ -22,14 +22,14 @@ kl_weight = 1e-7 -def KL_loss(z_mu, z_sigma): +def compute_kl_loss(z_mu, z_sigma): kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4]) return torch.sum(kl_loss) / kl_loss.shape[0] def generator_loss(gen_images, real_images, z_mu, z_sigma, disc_net, loss_perceptual): recons_loss = intensity_loss(gen_images, real_images) - kl_loss = KL_loss(z_mu, z_sigma) + kl_loss = compute_kl_loss(z_mu, z_sigma) p_loss = loss_perceptual(gen_images.float(), real_images.float()) loss_g = recons_loss + kl_weight * kl_loss + perceptual_weight * p_loss From 127421c7c424f1d95f9a0e3b92f68955df0320a2 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 20:13:28 -0700 Subject: [PATCH 102/130] update readme Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 4 ++-- models/brats_mri_generative_diffusion/docs/README.md | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 1f69d021..37f4a982 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -9,7 +9,7 @@ "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "@bundle_root", - "perceptual_loss_pretrained_path": "$@bundle_root +'/nv_resnet50.pth.tar'", + "perceptual_loss_model_weights_path": "$@bundle_root +'/nv_resnet50.pth.tar'", "download_brats": false, "train_batch_size": 2, "lr": 1e-05, @@ -64,7 +64,7 @@ "network_type": "resnet50", "is_fake_3d": true, "fake_3d_ratio": 0.2, - "pretrained_path": "@perceptual_loss_pretrained_path", + "pretrained_path": "@perceptual_loss_model_weights_path", "pretrained_state_dict_key": "state_dict" }, "dnetwork": "$@discriminator_def.to(@device)", diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index f5e3b365..d36c8bc6 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -42,6 +42,9 @@ The dataset can be downloaded automatically at the beginning of training. ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. +In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. +The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. + ### Training Configuration of Autoencoder The autoencoder was trained using the following configuration: From 3e54c1f72ce024bcd643c58e8b83f7f840ca4c75 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 12 May 2023 03:13:40 +0000 Subject: [PATCH 103/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d36c8bc6..2c067550 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -42,7 +42,7 @@ The dataset can be downloaded automatically at the beginning of training. ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. -In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. +In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. ### Training Configuration of Autoencoder From 1de2b1197b9fc1423b9a4392dff8c929aa775b46 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 20:24:14 -0700 Subject: [PATCH 104/130] add config file to inference autoencoder only Signed-off-by: Can-Zhao --- .../configs/inference_autoencoder.json | 136 ++++++++++++++++++ .../docs/README.md | 17 ++- 2 files changed, 150 insertions(+), 3 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/configs/inference_autoencoder.json diff --git a/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json new file mode 100644 index 00000000..19e8da21 --- /dev/null +++ b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json @@ -0,0 +1,136 @@ +{ + "imports": [ + "$import torch", + "$from datetime import datetime", + "$from pathlib import Path" + ], + "bundle_root": ".", + "model_dir": "$@bundle_root + '/models'", + "dataset_dir": "@bundle_root", + "output_dir": "$@bundle_root + '/output'", + "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", + "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", + "output_postfix": "", + "channel": 0, + "spacing": [ + 1.1, + 1.1, + 1.1 + ], + "spatial_dims": 3, + "image_channels": 1, + "latent_channels": 8, + "infer_patch_size": [ + 144, + 176, + 112 + ], + "autoencoder_def": { + "_target_": "generative.networks.nets.AutoencoderKL", + "spatial_dims": "@spatial_dims", + "in_channels": "@image_channels", + "out_channels": "@image_channels", + "latent_channels": "@latent_channels", + "num_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 + }, + "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", + "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", + "autoencoder": "$@autoencoder_def.to(@device)", + "preprocessing_transforms": [ + { + "_target_": "LoadImaged", + "keys": "image" + }, + { + "_target_": "EnsureChannelFirstd", + "keys": "image" + }, + { + "_target_": "Lambdad", + "keys": "image", + "func": "$lambda x: x[@channel, :, :, :]" + }, + { + "_target_": "AddChanneld", + "keys": "image" + }, + { + "_target_": "EnsureTyped", + "keys": "image" + }, + { + "_target_": "Orientationd", + "keys": "image", + "axcodes": "RAS" + }, + { + "_target_": "Spacingd", + "keys": "image", + "pixdim": "@spacing", + "mode": "bilinear" + } + ], + "crop_transforms": [ + { + "_target_": "CenterSpatialCropd", + "keys": "image", + "roi_size": "@infer_patch_size" + } + ], + "final_transforms": [ + { + "_target_": "ScaleIntensityRangePercentilesd", + "keys": "image", + "lower": 0, + "upper": 99.5, + "b_min": 0, + "b_max": 1 + } + ], + "preprocessing": { + "_target_": "Compose", + "transforms": "$@preprocessing_transforms + @crop_transforms + @final_transforms" + }, + "dataset": { + "_target_": "monai.apps.DecathlonDataset", + "root_dir": "@dataset_dir", + "task": "Task01_BrainTumour", + "section": "validation", + "cache_rate": 0.0, + "num_workers": 8, + "download": false, + "transform": "@preprocessing" + }, + "dataloader": { + "_target_": "DataLoader", + "dataset": "@dataset", + "batch_size": 1, + "shuffle": true, + "num_workers": 0 + }, + "saver": { + "_target_": "SaveImage", + "_requires_": "@create_output_dir", + "output_dir": "@output_dir", + "output_postfix": "@output_postfix" + }, + "input_img": "$monai.utils.first(@dataloader)['image'].to(@device)", + "recon_img": "$@autoencoder(@input_img)[0][0]", + "run": [ + "$@saver(@recon_img)" + ] +} diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index d36c8bc6..0d7ad18f 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -103,7 +103,9 @@ In addition to the Pythonic APIs, a few command line interfaces (CLI) are provid For more details usage instructions, visit the [MONAI Bundle Configuration Page](https://docs.monai.io/en/latest/config_syntax.html). -#### Execute Autoencoder Training (w/data download) +### Execute Autoencoder Training + +#### Execute Autoencoder Training (w/data download, single GPU) Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs using the BraTS dataset. If the dataset is not downloaded, it will be automatically downloaded and extracted to `./Task01_BrainTumour`. @@ -111,7 +113,7 @@ Assuming the current directory is the bundle directory, the following command wi python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True ``` -#### Execute Autoencoder Training +#### Execute Autoencoder Training (wo/data download, single GPU) If the dataset is already downloaded, make sure that `"dataset_dir"` in `configs/train_autoencoder.json` has the correct path to the dataset `Task01_BrainTumour`. Then, run: ``` @@ -125,7 +127,16 @@ To train with multiple GPUs, use the following command, which requires scaling u torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/multi_gpu_train_autoencoder.json']" --lr 8e-5 ``` -#### Execute Latent Diffusion Model Training +#### Check the Autoencoder Training result +The following code generates a reconstructed image from a random input image. +We can visualize it to see if the autoencoder is trained correctly. +``` +python -m monai.bundle run --config_file configs/inference_autoencoder.json +``` + +### Execute Latent Diffusion Training + +#### Execute Latent Diffusion Model Training (single GPU) After training the autoencoder, run the following command to train the latent diffusion model. This command will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. ``` From 58c799fc8937d77aa86482d362308e36f6440709 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 23:16:27 -0700 Subject: [PATCH 105/130] update infer Signed-off-by: Can-Zhao --- .../configs/inference_autoencoder.json | 21 +++++++++++++++---- .../configs/train_diffusion.json | 2 +- .../docs/README.md | 4 ++++ 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json index 19e8da21..501f342b 100644 --- a/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/inference_autoencoder.json @@ -10,7 +10,8 @@ "output_dir": "$@bundle_root + '/output'", "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", - "output_postfix": "", + "output_orig_postfix": "recon", + "output_recon_postfix": "orig", "channel": 0, "spacing": [ 1.1, @@ -122,15 +123,27 @@ "shuffle": true, "num_workers": 0 }, - "saver": { + "saver_orig": { "_target_": "SaveImage", "_requires_": "@create_output_dir", "output_dir": "@output_dir", - "output_postfix": "@output_postfix" + "output_postfix": "@output_orig_postfix", + "resample": false, + "padding_mode": "zeros" + }, + "saver_recon": { + "_target_": "SaveImage", + "_requires_": "@create_output_dir", + "output_dir": "@output_dir", + "output_postfix": "@output_recon_postfix", + "resample": false, + "padding_mode": "zeros" }, "input_img": "$monai.utils.first(@dataloader)['image'].to(@device)", "recon_img": "$@autoencoder(@input_img)[0][0]", "run": [ - "$@saver(@recon_img)" + "$@load_autoencoder", + "$@saver_orig(@input_img[0][0])", + "$@saver_recon(@recon_img)" ] } diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index b1be476b..a6d7f988 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -133,7 +133,7 @@ "trainer": { "_target_": "scripts.ldm_trainer.LDMTrainer", "device": "@device", - "max_epochs": 10000, + "max_epochs": 5000, "train_data_loader": "@train#dataloader", "network": "@diffusion", "autoencoder_model": "@autoencoder", diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 9db46a9a..03b81f70 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -134,6 +134,10 @@ We can visualize it to see if the autoencoder is trained correctly. python -m monai.bundle run --config_file configs/inference_autoencoder.json ``` +An example of reconstructed image from inference is shown below: +![Example reconstructed image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_recon_example.png) +If the autoencoder is trained correctly, the reconstructed image should look similar to original image. + ### Execute Latent Diffusion Training #### Execute Latent Diffusion Model Training (single GPU) From 9d6dc3e86ad1d5c5981ff9fa0e74c583aa944c81 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Thu, 11 May 2023 23:17:31 -0700 Subject: [PATCH 106/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 03b81f70..853c7d97 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -134,9 +134,10 @@ We can visualize it to see if the autoencoder is trained correctly. python -m monai.bundle run --config_file configs/inference_autoencoder.json ``` -An example of reconstructed image from inference is shown below: +An example of reconstructed image from inference is shown below. If the autoencoder is trained correctly, the reconstructed image should look similar to original image. + ![Example reconstructed image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_recon_example.png) -If the autoencoder is trained correctly, the reconstructed image should look similar to original image. + ### Execute Latent Diffusion Training From 42265dfe279d40e3ea99e611103a44879a9d45b2 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 12 May 2023 23:34:32 -0700 Subject: [PATCH 107/130] seprate data download out, update readme Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 1 - .../docs/README.md | 16 ++--- .../scripts/download_brats_data.py | 70 +++++++++++++++++++ 3 files changed, 76 insertions(+), 11 deletions(-) create mode 100644 models/brats_mri_generative_diffusion/scripts/download_brats_data.py diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 37f4a982..1d9c57ba 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -10,7 +10,6 @@ "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "@bundle_root", "perceptual_loss_model_weights_path": "$@bundle_root +'/nv_resnet50.pth.tar'", - "download_brats": false, "train_batch_size": 2, "lr": 1e-05, "train_patch_size": [ diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 853c7d97..14977518 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -37,7 +37,11 @@ The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medical - Modality: MRI - Size: 388 3D volumes (1 channel used) -The dataset can be downloaded automatically at the beginning of training. +The data can be automatically downloaded from [Medical Decathlon](http://medicaldecathlon.com/). By running the following command, the Brats data will be downloaded from [Medical Decathlon](http://medicaldecathlon.com/) and extracted to `$dataset_dir` in [./config/train_autoencoder.json](./config/train_autoencoder.json). You will see a subfolder `Task01_BrainTumour` under `$dataset_dir`. By default, you will see `./Task01_BrainTumour`. + +```bash +python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json +``` ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. @@ -105,15 +109,7 @@ For more details usage instructions, visit the [MONAI Bundle Configuration Page] ### Execute Autoencoder Training -#### Execute Autoencoder Training (w/data download, single GPU) - -Assuming the current directory is the bundle directory, the following command will train the autoencoder network for 1500 epochs using the BraTS dataset. If the dataset is not downloaded, it will be automatically downloaded and extracted to `./Task01_BrainTumour`. - -``` -python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir ./ --download_brats True -``` - -#### Execute Autoencoder Training (wo/data download, single GPU) +#### Execute Autoencoder Training on single GPU If the dataset is already downloaded, make sure that `"dataset_dir"` in `configs/train_autoencoder.json` has the correct path to the dataset `Task01_BrainTumour`. Then, run: ``` diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py new file mode 100644 index 00000000..61b010fb --- /dev/null +++ b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py @@ -0,0 +1,70 @@ +# 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 argparse +import json +import logging +import sys + +from monai.config import print_config +from monai.utils import set_determinism +from monai.apps import DecathlonDataset + + +def main(): + parser = argparse.ArgumentParser(description="PyTorch Object Detection Training") + parser.add_argument( + "-e", + "--environment-file", + default="./config/environment.json", + help="environment json file that stores environment path", + ) + args = parser.parse_args() + + # Step 0: configuration + print_config() + + env_dict = json.load(open(args.environment_file, "r")) + + for k, v in env_dict.items(): + setattr(args, k, v) + + set_determinism(42) + + # Step 1: set data loader + train_ds = DecathlonDataset( + root_dir=args.dataset_dir, + task="Task01_BrainTumour", + section="training", # validation + cache_rate=0.0, # you may need a few Gb of RAM... Set to 0 otherwise + num_workers=8, + download=True, # Set download to True if the dataset hasnt been downloaded yet + seed=0, + ) + val_ds = DecathlonDataset( + root_dir=args.data_base_dir, + task="Task01_BrainTumour", + section="validation", # validation + cache_rate=0.0, # you may need a few Gb of RAM... Set to 0 otherwise + num_workers=8, + download=True, # Set download to True if the dataset hasnt been downloaded yet + seed=0, + ) + + +if __name__ == "__main__": + logging.basicConfig( + stream=sys.stdout, + level=logging.INFO, + format="[%(asctime)s.%(msecs)03d][%(levelname)5s](%(name)s) - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + main() From 08cc6bcfe8d87907c2f6ba70809430d26284dbc3 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 12 May 2023 23:35:26 -0700 Subject: [PATCH 108/130] seprate data download out, update readme Signed-off-by: Can-Zhao --- .../scripts/download_brats_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py index 61b010fb..f5979e85 100644 --- a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py +++ b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py @@ -50,7 +50,7 @@ def main(): seed=0, ) val_ds = DecathlonDataset( - root_dir=args.data_base_dir, + root_dir=args.dataset_dir, task="Task01_BrainTumour", section="validation", # validation cache_rate=0.0, # you may need a few Gb of RAM... Set to 0 otherwise From 68e96788e1c6e1fd56af735d45468b5a48645e61 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 12 May 2023 23:39:21 -0700 Subject: [PATCH 109/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 14977518..dea97b1b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -132,7 +132,7 @@ python -m monai.bundle run --config_file configs/inference_autoencoder.json An example of reconstructed image from inference is shown below. If the autoencoder is trained correctly, the reconstructed image should look similar to original image. -![Example reconstructed image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_recon_example.png) +![Example reconstructed image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_recon_example.jpg) ### Execute Latent Diffusion Training From 7fb9f9ee29fedfbf2a98ef4f2a3b91e4c9a07e79 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 12 May 2023 23:40:56 -0700 Subject: [PATCH 110/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index dea97b1b..7830961e 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -116,7 +116,7 @@ If the dataset is already downloaded, make sure that `"dataset_dir"` in `configs python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -#### Override the `train` config to execute multi-GPU training for Autoencoder: +#### Override the `train` config to execute multi-GPU training for Autoencoder To train with multiple GPUs, use the following command, which requires scaling up the learning rate according to the number of GPUs. ``` @@ -137,27 +137,27 @@ An example of reconstructed image from inference is shown below. If the autoenco ### Execute Latent Diffusion Training -#### Execute Latent Diffusion Model Training (single GPU) +#### Execute Latent Diffusion Model Training on single GPU After training the autoencoder, run the following command to train the latent diffusion model. This command will print out the scale factor of the latent feature space. If your autoencoder is well trained, this value should be close to 1.0. ``` python -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json']" ``` -#### Override the `train` config to execute multi-GPU training for Latent Diffusion Model: +#### Override the `train` config to execute multi-GPU training for Latent Diffusion Model To train with multiple GPUs, use the following command, which requires scaling up the learning rate according to the number of GPUs. ``` torchrun --standalone --nnodes=1 --nproc_per_node=8 -m monai.bundle run --config_file "['configs/train_autoencoder.json','configs/train_diffusion.json','configs/multi_gpu_train_autoencoder.json','configs/multi_gpu_train_diffusion.json']" --lr 8e-5 ``` -#### Execute inference: +#### Execute inference The following code generates a synthetic image from a random sampled noise. ``` python -m monai.bundle run --config_file configs/inference.json ``` -#### Export checkpoint to TorchScript file: +#### Export checkpoint to TorchScript file The Autoencoder can be exported into a TorchScript file. From 3b6658be07870f5e3faece284778561f2f02e1b2 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Fri, 12 May 2023 23:47:47 -0700 Subject: [PATCH 111/130] reformat Signed-off-by: Can-Zhao --- .../scripts/download_brats_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py index f5979e85..4eb55b98 100644 --- a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py +++ b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py @@ -14,9 +14,9 @@ import logging import sys +from monai.apps import DecathlonDataset from monai.config import print_config from monai.utils import set_determinism -from monai.apps import DecathlonDataset def main(): From 612fe75c9a5c5d44f005a8a133d79f592792208c Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 14 May 2023 21:02:05 -0700 Subject: [PATCH 112/130] update readme and config for pre-trained issue Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 2 +- .../docs/README.md | 7 +++- .../scripts/download_brats_data.py | 38 ++++++++----------- 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 1d9c57ba..a68b3380 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -9,7 +9,7 @@ "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "@bundle_root", - "perceptual_loss_model_weights_path": "$@bundle_root +'/nv_resnet50.pth.tar'", + "perceptual_loss_model_weights_path": null, "train_batch_size": 2, "lr": 1e-05, "train_patch_size": [ diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 7830961e..04a73322 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,6 +7,9 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. +In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. +The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. + #### Example synthetic image An example result from inference is shown below: ![Example synthetic image](https://developer.download.nvidia.com/assets/Clara/Images/monai_brain_image_gen_ldm3d_example_generation.png) @@ -46,8 +49,8 @@ python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. -In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. -The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. +- `"train_batch_size"` is number of 3D volumes loaded in each batch. +- `"perceptual_loss_model_weights_path"` is the path to load pretrained weights for perceptual loss. This path MUST be provided to ensure correct training. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. ### Training Configuration of Autoencoder The autoencoder was trained using the following configuration: diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py index 4eb55b98..15d459ba 100644 --- a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py +++ b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py @@ -12,11 +12,13 @@ import argparse import json import logging +import os +import shutil import sys +import tempfile -from monai.apps import DecathlonDataset +from monai.apps import download_and_extract from monai.config import print_config -from monai.utils import set_determinism def main(): @@ -37,27 +39,19 @@ def main(): for k, v in env_dict.items(): setattr(args, k, v) - set_determinism(42) - # Step 1: set data loader - train_ds = DecathlonDataset( - root_dir=args.dataset_dir, - task="Task01_BrainTumour", - section="training", # validation - cache_rate=0.0, # you may need a few Gb of RAM... Set to 0 otherwise - num_workers=8, - download=True, # Set download to True if the dataset hasnt been downloaded yet - seed=0, - ) - val_ds = DecathlonDataset( - root_dir=args.dataset_dir, - task="Task01_BrainTumour", - section="validation", # validation - cache_rate=0.0, # you may need a few Gb of RAM... Set to 0 otherwise - num_workers=8, - download=True, # Set download to True if the dataset hasnt been downloaded yet - seed=0, - ) + directory = os.environ.get("MONAI_DATA_DIRECTORY") + root_dir = tempfile.mkdtemp() if directory is None else directory + + msd_task = "Task01_BrainTumour" + resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/" + msd_task + ".tar" + compressed_file = os.path.join(root_dir, msd_task + ".tar") + + os.makedirs(args.dataset_dir, exist_ok=True) + download_and_extract(resource, compressed_file, args.dataset_dir) + + if directory is None: + shutil.rmtree(root_dir) if __name__ == "__main__": From 11bb3f5aff7c357059a026eb24c3fbba6b0401e5 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 14 May 2023 21:04:19 -0700 Subject: [PATCH 113/130] update readme Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 04a73322..c8dcf016 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -8,7 +8,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. -The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. +The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. #### Example synthetic image An example result from inference is shown below: @@ -47,7 +47,7 @@ python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json ``` ## Training Configuration -If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](configs/train_autoencoder.json) and [configs/train_diffusion.json](configs/train_diffusion.json) configuration files. +If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json) and [configs/train_diffusion.json](../configs/train_diffusion.json) configuration files. - `"train_batch_size"` is number of 3D volumes loaded in each batch. - `"perceptual_loss_model_weights_path"` is the path to load pretrained weights for perceptual loss. This path MUST be provided to ensure correct training. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. From 5d44bf86dff23aff7991b8eb6c84740d463a719d Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 14 May 2023 21:05:31 -0700 Subject: [PATCH 114/130] update config for download Signed-off-by: Can-Zhao --- .../configs/train_autoencoder.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index a68b3380..2bb2bb73 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -143,7 +143,7 @@ "section": "training", "cache_rate": 1.0, "num_workers": 8, - "download": "@download_brats", + "download": false, "transform": "@train#preprocessing" }, "dataloader": { From de7d6e536da1599a0913b93cf3a56a817e98a654 Mon Sep 17 00:00:00 2001 From: Can-Zhao Date: Sun, 14 May 2023 21:51:38 -0700 Subject: [PATCH 115/130] typo Signed-off-by: Can-Zhao --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index c8dcf016..9ae08fbc 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -50,7 +50,7 @@ python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json) and [configs/train_diffusion.json](../configs/train_diffusion.json) configuration files. - `"train_batch_size"` is number of 3D volumes loaded in each batch. -- `"perceptual_loss_model_weights_path"` is the path to load pretrained weights for perceptual loss. This path MUST be provided to ensure correct training. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. +- `"perceptual_loss_model_weights_path"` is the path to load pre-trained weights for perceptual loss. This path MUST be provided to ensure correct training. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. ### Training Configuration of Autoencoder The autoencoder was trained using the following configuration: From a6677101198af9e1aca04f036f056e4b0abc4ddb Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 13:19:23 +0800 Subject: [PATCH 116/130] remove download flag Signed-off-by: Yiheng Wang --- .../brats_mri_generative_diffusion/configs/train_diffusion.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_diffusion.json b/models/brats_mri_generative_diffusion/configs/train_diffusion.json index a6d7f988..58ee2648 100644 --- a/models/brats_mri_generative_diffusion/configs/train_diffusion.json +++ b/models/brats_mri_generative_diffusion/configs/train_diffusion.json @@ -91,7 +91,7 @@ "section": "training", "cache_rate": 1.0, "num_workers": 8, - "download": "@download_brats", + "download": false, "transform": "@train#preprocessing" }, "dataloader": { From a84bca148b6d4ab145e51b19207e72ed82d4f2a5 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 13:24:50 +0800 Subject: [PATCH 117/130] fix typo Signed-off-by: Yiheng Wang --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 9ae08fbc..90afafa6 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -8,7 +8,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. -The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvison](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. +The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. #### Example synthetic image An example result from inference is shown below: From eddbfe3dca242485d8ed96a1ffba2cb6e619a249 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 14:01:16 +0800 Subject: [PATCH 118/130] reformat pretrained Signed-off-by: Yiheng Wang --- .../configs/train_autoencoder.json | 2 ++ models/brats_mri_generative_diffusion/docs/README.md | 9 ++++----- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 2bb2bb73..4f3b90ff 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -9,6 +9,7 @@ "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", "dataset_dir": "@bundle_root", + "pretrained": false, "perceptual_loss_model_weights_path": null, "train_batch_size": 2, "lr": 1e-05, @@ -63,6 +64,7 @@ "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" }, diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 90afafa6..ec39e898 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,8 +7,10 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -In this bundle, the autoencoder uses perceptual loss, which is based on a model with pre-trained weights from some internal data. This model is frozen and will not be trained in the bundle. -The path of the model is specified in `perceptual_loss_model_weights_path` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json). The [MONAI Generative Model repo](https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/losses/perceptual.py#L64-L69) and [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) also provide pre-trained weights but may be for non-commercial use only. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. +In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights. The network is forzen and will not be trained in the bundle. In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: +1. if set `pretrained` to `True`, ImageNet pretrained weights from [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) will be used. However, the weights are for non-commercial use only. +2. if set `pretrained` to `True` and specifies the `perceptual_loss_model_weights_path` parameter, users are able to load weights from a local path. This is the way this bundle used to train, and the pre-trained weights are from some internal data. +Please note that each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. #### Example synthetic image An example result from inference is shown below: @@ -49,9 +51,6 @@ python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json) and [configs/train_diffusion.json](../configs/train_diffusion.json) configuration files. -- `"train_batch_size"` is number of 3D volumes loaded in each batch. -- `"perceptual_loss_model_weights_path"` is the path to load pre-trained weights for perceptual loss. This path MUST be provided to ensure correct training. Each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. - ### Training Configuration of Autoencoder The autoencoder was trained using the following configuration: From 386a19970a6c9f3f8f73e9edf44a1844153d317f Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 14:03:08 +0800 Subject: [PATCH 119/130] adjust sentence Signed-off-by: Yiheng Wang --- models/brats_mri_generative_diffusion/docs/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index ec39e898..dcfab28d 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -10,6 +10,7 @@ This model is a generator for creating images like the Flair MRIs based on BraTS In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights. The network is forzen and will not be trained in the bundle. In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: 1. if set `pretrained` to `True`, ImageNet pretrained weights from [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) will be used. However, the weights are for non-commercial use only. 2. if set `pretrained` to `True` and specifies the `perceptual_loss_model_weights_path` parameter, users are able to load weights from a local path. This is the way this bundle used to train, and the pre-trained weights are from some internal data. + Please note that each user is responsible for checking the data source of the pre-trained models, the applicable licenses, and determining if suitable for the intended use. #### Example synthetic image From 9f03dd10057ecbb346e7c15c3759099ee0651671 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 14:46:40 +0800 Subject: [PATCH 120/130] use correct install commit id Signed-off-by: Yiheng Wang --- .../install_brats_mri_generative_diffusion_dependency.sh | 2 +- models/brats_mri_generative_diffusion/docs/README.md | 7 +------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index d9f6bd48..f808a82d 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -11,6 +11,6 @@ git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ -git checkout ad973b8157130eaa146f35628f96eebd2b9d207d +git checkout f969c24f88d013dc0045fb7b2885a01fb219992b python setup.py install cd .. diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index dcfab28d..b730d576 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -23,14 +23,9 @@ An example result from inference is shown below: [MONAI generative models](https://github.com/Project-MONAI/GenerativeModels) can be installed by ``` pip install lpips==0.1.4 -pip install git+https://github.com/Project-MONAI/GenerativeModels.git@0.2.0 -``` - -Or we can install it from source -``` -pip install lpips==0.1.4 git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ +git checkout f969c24f88d013dc0045fb7b2885a01fb219992b python setup.py install cd .. ``` From 1504191202fa84534bd496930dbb191543c39f0c Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 15:00:54 +0800 Subject: [PATCH 121/130] rephrase sentence Signed-off-by: Yiheng Wang --- models/brats_mri_generative_diffusion/docs/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index b730d576..48172841 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights. The network is forzen and will not be trained in the bundle. In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: +In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights (The network is frozen and will not be trained in the bundle). In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: 1. if set `pretrained` to `True`, ImageNet pretrained weights from [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) will be used. However, the weights are for non-commercial use only. 2. if set `pretrained` to `True` and specifies the `perceptual_loss_model_weights_path` parameter, users are able to load weights from a local path. This is the way this bundle used to train, and the pre-trained weights are from some internal data. From cecb03ac5a4a22fab5c3efa4b509c49551357da4 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 15:11:35 +0800 Subject: [PATCH 122/130] fix download related errors Signed-off-by: Yiheng Wang --- models/brats_mri_generative_diffusion/docs/README.md | 4 ++-- .../scripts/download_brats_data.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 48172841..97ac2f24 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -38,10 +38,10 @@ The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medical - Modality: MRI - Size: 388 3D volumes (1 channel used) -The data can be automatically downloaded from [Medical Decathlon](http://medicaldecathlon.com/). By running the following command, the Brats data will be downloaded from [Medical Decathlon](http://medicaldecathlon.com/) and extracted to `$dataset_dir` in [./config/train_autoencoder.json](./config/train_autoencoder.json). You will see a subfolder `Task01_BrainTumour` under `$dataset_dir`. By default, you will see `./Task01_BrainTumour`. +The data can be automatically downloaded from [Medical Decathlon](http://medicaldecathlon.com/). By running the following command, the Brats data will be downloaded from [Medical Decathlon](http://medicaldecathlon.com/) and extracted to `$dataset_dir` in [./configs/train_autoencoder.json](./configs/train_autoencoder.json). You will see a subfolder `Task01_BrainTumour` under `$dataset_dir`. By default, you will see `./Task01_BrainTumour`. ```bash -python ./scripts/download_brats_data.py -e ./config/train_autoencoder.json +python ./scripts/download_brats_data.py -e ./configs/train_autoencoder.json ``` ## Training Configuration diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py index 15d459ba..a3780eef 100644 --- a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py +++ b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py @@ -26,7 +26,7 @@ def main(): parser.add_argument( "-e", "--environment-file", - default="./config/environment.json", + default="./configs/train_autoencoder.json", help="environment json file that stores environment path", ) args = parser.parse_args() From a0bab96fdd042d6eaf853d9ed77d1a7dbfe303de Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 16:00:34 +0800 Subject: [PATCH 123/130] remove error file and update data part of readme Signed-off-by: Yiheng Wang --- .../configs/train_autoencoder.json | 2 +- .../docs/README.md | 15 ++--- .../scripts/download_brats_data.py | 64 ------------------- 3 files changed, 8 insertions(+), 73 deletions(-) delete mode 100644 models/brats_mri_generative_diffusion/scripts/download_brats_data.py diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index 4f3b90ff..f6efce12 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -8,7 +8,7 @@ "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", "ckpt_dir": "$@bundle_root + '/models'", "tf_dir": "$@bundle_root + '/eval'", - "dataset_dir": "@bundle_root", + "dataset_dir": "/workspace/data/medical/Task01_BrainTumour", "pretrained": false, "perceptual_loss_model_weights_path": null, "train_batch_size": 2, diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 97ac2f24..6941b054 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -31,19 +31,13 @@ cd .. ``` ## Data -The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). +The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). Users can find more details on the datasets (`Task01_BrainTumour`) at http://medicaldecathlon.com/. - Target: Image Generation - Task: Synthesis - Modality: MRI - Size: 388 3D volumes (1 channel used) -The data can be automatically downloaded from [Medical Decathlon](http://medicaldecathlon.com/). By running the following command, the Brats data will be downloaded from [Medical Decathlon](http://medicaldecathlon.com/) and extracted to `$dataset_dir` in [./configs/train_autoencoder.json](./configs/train_autoencoder.json). You will see a subfolder `Task01_BrainTumour` under `$dataset_dir`. By default, you will see `./Task01_BrainTumour`. - -```bash -python ./scripts/download_brats_data.py -e ./configs/train_autoencoder.json -``` - ## Training Configuration If you have a GPU with less than 32G of memory, you may need to decrease the batch size when training. To do so, modify the `train_batch_size` parameter in the [configs/train_autoencoder.json](../configs/train_autoencoder.json) and [configs/train_diffusion.json](../configs/train_diffusion.json) configuration files. @@ -108,12 +102,17 @@ For more details usage instructions, visit the [MONAI Bundle Configuration Page] ### Execute Autoencoder Training #### Execute Autoencoder Training on single GPU -If the dataset is already downloaded, make sure that `"dataset_dir"` in `configs/train_autoencoder.json` has the correct path to the dataset `Task01_BrainTumour`. Then, run: ``` python -m monai.bundle run --config_file configs/train_autoencoder.json ``` +Please note that if the default dataset path is not modified with the actual path in the bundle config files, you can also override it by using `--dataset_dir`: + +``` +python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir +``` + #### Override the `train` config to execute multi-GPU training for Autoencoder To train with multiple GPUs, use the following command, which requires scaling up the learning rate according to the number of GPUs. diff --git a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py b/models/brats_mri_generative_diffusion/scripts/download_brats_data.py deleted file mode 100644 index a3780eef..00000000 --- a/models/brats_mri_generative_diffusion/scripts/download_brats_data.py +++ /dev/null @@ -1,64 +0,0 @@ -# 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 argparse -import json -import logging -import os -import shutil -import sys -import tempfile - -from monai.apps import download_and_extract -from monai.config import print_config - - -def main(): - parser = argparse.ArgumentParser(description="PyTorch Object Detection Training") - parser.add_argument( - "-e", - "--environment-file", - default="./configs/train_autoencoder.json", - help="environment json file that stores environment path", - ) - args = parser.parse_args() - - # Step 0: configuration - print_config() - - env_dict = json.load(open(args.environment_file, "r")) - - for k, v in env_dict.items(): - setattr(args, k, v) - - # Step 1: set data loader - directory = os.environ.get("MONAI_DATA_DIRECTORY") - root_dir = tempfile.mkdtemp() if directory is None else directory - - msd_task = "Task01_BrainTumour" - resource = "https://msd-for-monai.s3-us-west-2.amazonaws.com/" + msd_task + ".tar" - compressed_file = os.path.join(root_dir, msd_task + ".tar") - - os.makedirs(args.dataset_dir, exist_ok=True) - download_and_extract(resource, compressed_file, args.dataset_dir) - - if directory is None: - shutil.rmtree(root_dir) - - -if __name__ == "__main__": - logging.basicConfig( - stream=sys.stdout, - level=logging.INFO, - format="[%(asctime)s.%(msecs)03d][%(levelname)5s](%(name)s) - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - main() From a08d963b084f3c34d8daa8223cb7292f2a157199 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 16:21:45 +0800 Subject: [PATCH 124/130] remove duplicate link and modify default dataset dir Signed-off-by: Yiheng Wang --- .../configs/train_autoencoder.json | 2 +- models/brats_mri_generative_diffusion/docs/README.md | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json index f6efce12..28136f51 100644 --- a/models/brats_mri_generative_diffusion/configs/train_autoencoder.json +++ b/models/brats_mri_generative_diffusion/configs/train_autoencoder.json @@ -8,7 +8,7 @@ "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/Task01_BrainTumour", + "dataset_dir": "/workspace/data/medical", "pretrained": false, "perceptual_loss_model_weights_path": null, "train_batch_size": 2, diff --git a/models/brats_mri_generative_diffusion/docs/README.md b/models/brats_mri_generative_diffusion/docs/README.md index 6941b054..f081ca3b 100644 --- a/models/brats_mri_generative_diffusion/docs/README.md +++ b/models/brats_mri_generative_diffusion/docs/README.md @@ -7,7 +7,7 @@ This model is trained on BraTS 2016 and 2017 data from [Medical Decathlon](http: This model is a generator for creating images like the Flair MRIs based on BraTS 2016 and 2017 data. It was trained as a 3d latent diffusion model and accepts Gaussian random noise as inputs to produce an image output. The `train_autoencoder.json` file describes the training process of the variational autoencoder with GAN loss. The `train_diffusion.json` file describes the training process of the 3D latent diffusion model. -In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights (The network is frozen and will not be trained in the bundle). In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: +In this bundle, the autoencoder uses perceptual loss, which is based on ResNet50 with pre-trained weights (the network is frozen and will not be trained in the bundle). In default, the `pretrained` parameter is specified as `False` in `train_autoencoder.json`. To ensure correct training, changing the default settings is necessary. There are two ways to utilize pretrained weights: 1. if set `pretrained` to `True`, ImageNet pretrained weights from [torchvision](https://pytorch.org/vision/stable/_modules/torchvision/models/resnet.html#ResNet50_Weights) will be used. However, the weights are for non-commercial use only. 2. if set `pretrained` to `True` and specifies the `perceptual_loss_model_weights_path` parameter, users are able to load weights from a local path. This is the way this bundle used to train, and the pre-trained weights are from some internal data. @@ -31,7 +31,7 @@ cd .. ``` ## Data -The training data is BraTS 2016 and 2017 from [Medical Decathlon](http://medicaldecathlon.com/). Users can find more details on the datasets (`Task01_BrainTumour`) at http://medicaldecathlon.com/. +The training data is BraTS 2016 and 2017 from the Medical Segmentation Decathalon. Users can find more details on the dataset (`Task01_BrainTumour`) at http://medicaldecathlon.com/. - Target: Image Generation - Task: Synthesis @@ -107,7 +107,7 @@ For more details usage instructions, visit the [MONAI Bundle Configuration Page] python -m monai.bundle run --config_file configs/train_autoencoder.json ``` -Please note that if the default dataset path is not modified with the actual path in the bundle config files, you can also override it by using `--dataset_dir`: +Please note that if the default dataset path is not modified with the actual path (it should be the path that contains `Task01_BrainTumour`) in the bundle config files, you can also override it by using `--dataset_dir`: ``` python -m monai.bundle run --config_file configs/train_autoencoder.json --dataset_dir From 1f50f395bce5419721f6c804dd19a756a9c559aa Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 19:20:35 +0800 Subject: [PATCH 125/130] export installed path in ci Signed-off-by: Yiheng Wang --- .../install_brats_mri_generative_diffusion_dependency.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index f808a82d..e0555a23 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -14,3 +14,4 @@ cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b python setup.py install cd .. +export PYTHONPATH=GenerativeModels/:$PYTHONPATH From ea01892e82e650c358541413623adb4692b9d9d5 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 20:36:29 +0800 Subject: [PATCH 126/130] rewrite install lib ci Signed-off-by: Yiheng Wang --- ci/get_bundle_requirements.py | 16 +++++++++++++++- ..._brats_mri_generative_diffusion_dependency.sh | 1 - ci/run_premerge_gpu.sh | 6 ++++++ ci/verify_bundle.py | 8 -------- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/ci/get_bundle_requirements.py b/ci/get_bundle_requirements.py index 03dddf06..e9de7a2e 100644 --- a/ci/get_bundle_requirements.py +++ b/ci/get_bundle_requirements.py @@ -13,6 +13,7 @@ import argparse import os +from bundle_custom_data import install_dependency_dict from utils import get_json_dict @@ -44,11 +45,24 @@ def get_requirements(bundle, models_path): print(requirements_file_name) +def get_install_script(bundle): + # install extra dependencies if needed + script_path = "" + if bundle in install_dependency_dict.keys(): + script_path = install_dependency_dict[bundle] + print(script_path) + + if __name__ == "__main__": parser = argparse.ArgumentParser(description="") parser.add_argument("-b", "--b", type=str, help="bundle name.") parser.add_argument("-p", "--p", type=str, default="models", help="models path.") + parser.add_argument("--get_script", type=bool, default=False, help="whether to get the install script.") args = parser.parse_args() bundle = args.b models_path = args.p - get_requirements(bundle, models_path) + get_script = args.get_script + if get_script is True: + get_install_script(bundle) + else: + get_requirements(bundle, models_path) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index e0555a23..f808a82d 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -14,4 +14,3 @@ cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b python setup.py install cd .. -export PYTHONPATH=GenerativeModels/:$PYTHONPATH diff --git a/ci/run_premerge_gpu.sh b/ci/run_premerge_gpu.sh index 4a7abcc5..e3240348 100755 --- a/ci/run_premerge_gpu.sh +++ b/ci/run_premerge_gpu.sh @@ -65,6 +65,12 @@ verify_bundle() { echo "install required libraries for bundle: $bundle" pipenv install -r "$requirements" fi + # get extra install script if exists + extra_script=$(pipenv run python $(pwd)/ci/get_bundle_requirements.py --b "$bundle" --get_script True) + if [ ! -z "$extra_script" ]; then + echo "install extra libraries with script: $extra_script" + bash $extra_script + fi # verify bundle pipenv run python $(pwd)/ci/verify_bundle.py --b "$bundle" remove_pipenv diff --git a/ci/verify_bundle.py b/ci/verify_bundle.py index f6b11c2a..b48f800b 100644 --- a/ci/verify_bundle.py +++ b/ci/verify_bundle.py @@ -11,7 +11,6 @@ import argparse import os -import subprocess import sys from typing import List @@ -20,7 +19,6 @@ exclude_verify_preferred_files_list, exclude_verify_shape_list, exclude_verify_torchscript_list, - install_dependency_dict, ) from monai.bundle import ckpt_export, verify_metadata, verify_net_in_out from monai.bundle.config_parser import ConfigParser @@ -310,12 +308,6 @@ def verify_torchscript( def verify(bundle, models_path="models", mode="full"): print(f"start verifying {bundle}:") - # install extra dependencies if needed - if bundle in install_dependency_dict.keys(): - script_path = install_dependency_dict[bundle] - install_cmd = f"bash {script_path}" - call_status = subprocess.run(install_cmd, shell=True) - call_status.check_returncode() # add bundle path to ensure custom code can be used sys.path = [os.path.join(models_path, bundle)] + sys.path # verify bundle directory From d31e59dfd69aeff757fa4e5887b88d31c0c4caa8 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 21:00:58 +0800 Subject: [PATCH 127/130] add cython install Signed-off-by: Yiheng Wang --- .../install_brats_mri_generative_diffusion_dependency.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index f808a82d..a9cf2178 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -9,6 +9,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# install cython to avoid error: +# https://github.com/Project-MONAI/model-zoo/actions/runs/4980530305/jobs/8913560199#step:2:550 +pip install cython==0.29.34 git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b From bb287264719319f001f17c5a7faadcbe4eed8b37 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 22:09:09 +0800 Subject: [PATCH 128/130] use export pythonpath instead of setup install Signed-off-by: Yiheng Wang --- .../install_brats_mri_generative_diffusion_dependency.sh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index a9cf2178..d79fa8fe 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -11,9 +11,8 @@ # install cython to avoid error: # https://github.com/Project-MONAI/model-zoo/actions/runs/4980530305/jobs/8913560199#step:2:550 -pip install cython==0.29.34 git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b -python setup.py install cd .. +export PYTHONPATH=GenerativeModels/:$PYTHONPATH \ No newline at end of file From d3b8233c59fa795745012d79daeee95d60d00b00 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 May 2023 14:09:20 +0000 Subject: [PATCH 129/130] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../install_brats_mri_generative_diffusion_dependency.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index d79fa8fe..6dad07bc 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -15,4 +15,4 @@ git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b cd .. -export PYTHONPATH=GenerativeModels/:$PYTHONPATH \ No newline at end of file +export PYTHONPATH=GenerativeModels/:$PYTHONPATH From 77644acfb74cde5e1c766d17fbdeb99f1031fe43 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Mon, 15 May 2023 22:26:06 +0800 Subject: [PATCH 130/130] exclude ts test Signed-off-by: Yiheng Wang --- ci/bundle_custom_data.py | 1 + .../install_brats_mri_generative_diffusion_dependency.sh | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/ci/bundle_custom_data.py b/ci/bundle_custom_data.py index 20923eab..dbf91530 100644 --- a/ci/bundle_custom_data.py +++ b/ci/bundle_custom_data.py @@ -33,6 +33,7 @@ "wholeBrainSeg_Large_UNEST_segmentation", "breast_density_classification", "mednist_reg", + "brats_mri_generative_diffusion", ] # This dict is used for our CI tests to install required dependencies that cannot be installed by `pip install` directly. diff --git a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh index d79fa8fe..13207a1f 100644 --- a/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh +++ b/ci/install_scripts/install_brats_mri_generative_diffusion_dependency.sh @@ -14,5 +14,5 @@ git clone https://github.com/Project-MONAI/GenerativeModels.git cd GenerativeModels/ git checkout f969c24f88d013dc0045fb7b2885a01fb219992b +python setup.py install cd .. -export PYTHONPATH=GenerativeModels/:$PYTHONPATH \ No newline at end of file