Skip to content

Commit

Permalink
TRT support for MAISI (#701)
Browse files Browse the repository at this point in the history
### Description
TRT optimization support for MAISI. 
Depends on Project-MONAI/MONAI#8153

To successfully export MAISI, either 24.11 PyTorch container (or latest
public Torch nightly) is needed, or this patch needs to be applied to
24.10-based container:

```
--- /usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset14.bak     2024-10-09 01:38:04.920316673 +0000                                                   
+++ /usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset14.py      2024-10-09 01:38:25.228053951 +0000                                                   
@@ -148,7 +148,6 @@                                                                                                                                                   
         is_causal and symbolic_helper._is_none(attn_mask)                                                                                                            
     ), "is_causal and attn_mask cannot be set at the same time"                                                                                                      
                                                                                                                                                                      
-    scale = symbolic_helper._maybe_get_const(scale, "f")                                                                                                             
     if symbolic_helper._is_none(scale):                                                                                                                              
         scale = _attention_scale(g, query)

---------

Signed-off-by: Boris Fomitchev <[email protected]>
Signed-off-by: Yiheng Wang <[email protected]>
Co-authored-by: Yiheng Wang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
3 people authored Dec 12, 2024
1 parent d486306 commit 397f35c
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 1 deletion.
19 changes: 19 additions & 0 deletions models/maisi_ct_generative/configs/inference_trt.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
{
"+imports": [
"$from monai.networks import trt_compile"
],
"c_trt_args": {
"export_args": {
"dynamo": "$False",
"report": "$True"
},
"output_lists": [
[
-1
],
[]
]
},
"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)"
}
3 changes: 2 additions & 1 deletion models/maisi_ct_generative/configs/metadata.json
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
{
"schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20240318.json",
"version": "0.4.5",
"version": "0.4.6",
"changelog": {
"0.4.6": "add TensorRT support",
"0.4.5": "update README",
"0.4.4": "update issue for IgniteInfo",
"0.4.3": "remove download large files, add weights_only when loading weights and add label_dict to large files",
Expand Down

0 comments on commit 397f35c

Please sign in to comment.