From 04a932695a877de98f5b84515857baa508703217 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Thu, 19 Dec 2024 18:01:50 -0800 Subject: [PATCH 1/3] Added TRT config for inference Signed-off-by: Boris Fomitchev --- generation/maisi/configs/config_infer.json | 7 ++++- generation/maisi/configs/config_trt.json | 22 ++++++++++++++++ generation/maisi/scripts/inference.py | 30 +++++++++++++++++----- 3 files changed, 51 insertions(+), 8 deletions(-) create mode 100644 generation/maisi/configs/config_trt.json diff --git a/generation/maisi/configs/config_infer.json b/generation/maisi/configs/config_infer.json index 9586081af7..fc08a7bda4 100644 --- a/generation/maisi/configs/config_infer.json +++ b/generation/maisi/configs/config_infer.json @@ -18,5 +18,10 @@ 2.0 ], "autoencoder_sliding_window_infer_size": [48,48,48], - "autoencoder_sliding_window_infer_overlap": 0.25 + "autoencoder_sliding_window_infer_overlap": 0.25, + "controlnet": "$@controlnet_def", + "diffusion_unet": "$@diffusion_unet_def", + "autoencoder": "$@autoencoder_def", + "mask_generation_autoencoder": "$@mask_generation_autoencoder_def", + "mask_generation_diffusion": "$@mask_generation_diffusion_def" } diff --git a/generation/maisi/configs/config_trt.json b/generation/maisi/configs/config_trt.json new file mode 100644 index 0000000000..fc52486e39 --- /dev/null +++ b/generation/maisi/configs/config_trt.json @@ -0,0 +1,22 @@ +{ + "+imports": [ + "$from monai.networks import trt_compile" + ], + "c_trt_args": { + "export_args": { + "dynamo": "$False", + "report": "$True" + }, + "output_lists": [ + [ + -1 + ], + [ + ] + ] + }, + "device": "cuda", + "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path, @c_trt_args)", + "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)", + "mask_generation_diffusion": "$trt_compile(@mask_generation_diffusion_def.to(@device), @trained_mask_generation_diffusion_path)" +} diff --git a/generation/maisi/scripts/inference.py b/generation/maisi/scripts/inference.py index 8220f200cc..6fa85e7b36 100644 --- a/generation/maisi/scripts/inference.py +++ b/generation/maisi/scripts/inference.py @@ -48,6 +48,12 @@ def main(): default="./configs/config_infer.json", help="config json file that stores inference hyper-parameters", ) + parser.add_argument( + "-x", + "--extra-config-file", + default=None, + help="config json file that stores inference extra parameters", + ) parser.add_argument( "-s", "--random-seed", @@ -122,7 +128,7 @@ def main(): print(f"{k}: {val}") print("Global config variables have been loaded.") - # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc. + # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc. # # The information for the inference input, like body region and anatomy to generate, is stored in "./configs/config_infer.json". # Please refer to README.md for the details. @@ -135,11 +141,21 @@ def main(): # override num_split if asked if "autoencoder_tp_num_splits" in config_infer_dict: args.autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] - args.mask_generation_autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] + args.mask_generation_autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] for k, v in config_infer_dict.items(): setattr(args, k, v) print(f"{k}: {v}") + # + # ## Read in optional extra configuration setting - typically acceleration options (TRT) + # + # + if args.extra_config_file is not None: + extra_config_dict = json.load(open(args.extra_config_file, "r")) + for k, v in extra_config_dict.items(): + setattr(args, k, v) + print(f"{k}: {v}") + check_input( args.body_region, args.anatomy_list, @@ -158,25 +174,25 @@ def main(): device = torch.device("cuda") - autoencoder = define_instance(args, "autoencoder_def").to(device) + autoencoder = define_instance(args, "autoencoder").to(device) checkpoint_autoencoder = torch.load(args.trained_autoencoder_path) autoencoder.load_state_dict(checkpoint_autoencoder) - diffusion_unet = define_instance(args, "diffusion_unet_def").to(device) + diffusion_unet = define_instance(args, "diffusion_unet").to(device) checkpoint_diffusion_unet = torch.load(args.trained_diffusion_path) diffusion_unet.load_state_dict(checkpoint_diffusion_unet["unet_state_dict"], strict=True) scale_factor = checkpoint_diffusion_unet["scale_factor"].to(device) - controlnet = define_instance(args, "controlnet_def").to(device) + controlnet = define_instance(args, "controlnet").to(device) checkpoint_controlnet = torch.load(args.trained_controlnet_path) monai.networks.utils.copy_model_state(controlnet, diffusion_unet.state_dict()) controlnet.load_state_dict(checkpoint_controlnet["controlnet_state_dict"], strict=True) - mask_generation_autoencoder = define_instance(args, "mask_generation_autoencoder_def").to(device) + mask_generation_autoencoder = define_instance(args, "mask_generation_autoencoder").to(device) checkpoint_mask_generation_autoencoder = torch.load(args.trained_mask_generation_autoencoder_path) mask_generation_autoencoder.load_state_dict(checkpoint_mask_generation_autoencoder) - mask_generation_diffusion_unet = define_instance(args, "mask_generation_diffusion_def").to(device) + mask_generation_diffusion_unet = define_instance(args, "mask_generation_diffusion").to(device) checkpoint_mask_generation_diffusion_unet = torch.load(args.trained_mask_generation_diffusion_path) mask_generation_diffusion_unet.load_state_dict(checkpoint_mask_generation_diffusion_unet["unet_state_dict"]) mask_generation_scale_factor = checkpoint_mask_generation_diffusion_unet["scale_factor"] From 0b4545bdd47186e03fb9bc4d0e74ca89922aa5a6 Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Thu, 19 Dec 2024 18:11:22 -0800 Subject: [PATCH 2/3] cleanup Signed-off-by: Boris Fomitchev --- generation/maisi/scripts/inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/generation/maisi/scripts/inference.py b/generation/maisi/scripts/inference.py index 6fa85e7b36..968d5bf49f 100644 --- a/generation/maisi/scripts/inference.py +++ b/generation/maisi/scripts/inference.py @@ -128,7 +128,7 @@ def main(): print(f"{k}: {val}") print("Global config variables have been loaded.") - # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc. + # ## Read in configuration setting, including network definition, body region and anatomy to generate, etc. # # The information for the inference input, like body region and anatomy to generate, is stored in "./configs/config_infer.json". # Please refer to README.md for the details. @@ -141,7 +141,7 @@ def main(): # override num_split if asked if "autoencoder_tp_num_splits" in config_infer_dict: args.autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] - args.mask_generation_autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] + args.mask_generation_autoencoder_def["num_splits"] = config_infer_dict["autoencoder_tp_num_splits"] for k, v in config_infer_dict.items(): setattr(args, k, v) print(f"{k}: {v}") From 4f950736322120d103c04505407384291d98b97d Mon Sep 17 00:00:00 2001 From: Boris Fomitchev Date: Wed, 15 Jan 2025 13:18:14 -0800 Subject: [PATCH 3/3] Added README chapter Signed-off-by: Boris Fomitchev --- generation/maisi/README.md | 10 ++++++++++ generation/maisi/configs/config_trt.json | 2 ++ 2 files changed, 12 insertions(+) diff --git a/generation/maisi/README.md b/generation/maisi/README.md index e7e4e9e477..f51cdf61fc 100644 --- a/generation/maisi/README.md +++ b/generation/maisi/README.md @@ -172,6 +172,16 @@ python -m scripts.inference -c ./configs/config_maisi.json -i ./configs/config_i Please refer to [maisi_inference_tutorial.ipynb](maisi_inference_tutorial.ipynb) for the tutorial for MAISI model inference. + +#### Accelerated Inference with TensorRT: +To run the inference script with TensorRT acceleration, please run: +```bash +export MONAI_DATA_DIRECTORY= +python -m scripts.inference -c ./configs/config_maisi.json -i ./configs/config_infer.json -e ./configs/environment.json -x ./configs/config_trt.json --random-seed 0 +``` +Extra config file, [./configs/config_trt.json](./configs/config_trt.json) is using `trt_compile()` utility from MONAI to convert select modules to TensorRT by overriding their definitions from [./configs/config_infer.json](./configs/config_infer.json). + + #### Quality Check: We have implemented a quality check function for the generated CT images. The main idea behind this function is to ensure that the Hounsfield units (HU) intensity for each organ in the CT images remains within a defined range. For each training image used in the Diffusion network, we computed the median value for a few major organs. Then we summarize the statistics of these median values and save it to [./configs/image_median_statistics.json](./configs/image_median_statistics.json). During inference, for each generated image, we compute the median HU values for the major organs and check whether they fall within the normal range. diff --git a/generation/maisi/configs/config_trt.json b/generation/maisi/configs/config_trt.json index fc52486e39..de6469fe0d 100644 --- a/generation/maisi/configs/config_trt.json +++ b/generation/maisi/configs/config_trt.json @@ -18,5 +18,7 @@ "device": "cuda", "controlnet": "$trt_compile(@controlnet_def.to(@device), @trained_controlnet_path, @c_trt_args)", "diffusion_unet": "$trt_compile(@diffusion_unet_def.to(@device), @trained_diffusion_path)", + "autoencoder": "$trt_compile(@autoencoder_def.to(@device), @trained_autoencoder_path, submodule='decoder')", + "mask_generation_autoencoder": "$trt_compile(@mask_generation_autoencoder_def.to(@device), @trained_mask_generation_autoencoder_path, submodule='decoder')", "mask_generation_diffusion": "$trt_compile(@mask_generation_diffusion_def.to(@device), @trained_mask_generation_diffusion_path)" }