Fixed the issue in comments

Author: cehongwang

examples/dynamo/mutable_torchtrt_module_example.py

@@ -14,6 +14,7 @@
 1. Sample workflow of Mutable Torch TensorRT Module with ResNet 18
 2. Save a Mutable Torch TensorRT Module
 3. Integration with Huggingface pipeline in LoRA use case
+4. Usage of dynamic shape with Mutable Torch TensorRT Module
 """
 
 import numpy as np
@@ -25,89 +26,123 @@
 torch.manual_seed(5)
 inputs = [torch.rand((1, 3, 224, 224)).to("cuda")]
 
-# %%
-# Initialize the Mutable Torch TensorRT Module with settings.
-# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-settings = {
-    "use_python": False,
-    "enabled_precisions": {torch.float32},
-    "immutable_weights": False,
-}
+# # %%
+# # Initialize the Mutable Torch TensorRT Module with settings.
+# # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# settings = {
+#     "use_python": False,
+#     "enabled_precisions": {torch.float32},
+#     "immutable_weights": False,
+# }
 
-model = models.resnet18(pretrained=True).eval().to("cuda")
-mutable_module = torch_trt.MutableTorchTensorRTModule(model, **settings)
-# You can use the mutable module just like the original pytorch module. The compilation happens while you first call the mutable module.
-mutable_module(*inputs)
+# model = models.resnet18(pretrained=True).eval().to("cuda")
+# mutable_module = torch_trt.MutableTorchTensorRTModule(model, **settings)
+# # You can use the mutable module just like the original pytorch module. The compilation happens while you first call the mutable module.
+# mutable_module(*inputs)
 
-# %%
-# Make modifications to the mutable module.
-# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# # %%
+# # Make modifications to the mutable module.
+# # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-# %%
-# Making changes to mutable module can trigger refit or re-compilation. For example, loading a different state_dict and setting new weight values will trigger refit, and adding a module to the model will trigger re-compilation.
-model2 = models.resnet18(pretrained=False).eval().to("cuda")
-mutable_module.load_state_dict(model2.state_dict())
+# # %%
+# # Making changes to mutable module can trigger refit or re-compilation. For example, loading a different state_dict and setting new weight values will trigger refit, and adding a module to the model will trigger re-compilation.
+# model2 = models.resnet18(pretrained=False).eval().to("cuda")
+# mutable_module.load_state_dict(model2.state_dict())
 
 
-# Check the output
-# The refit happens while you call the mutable module again.
-expected_outputs, refitted_outputs = model2(*inputs), mutable_module(*inputs)
-for expected_output, refitted_output in zip(expected_outputs, refitted_outputs):
-    assert torch.allclose(
-        expected_output, refitted_output, 1e-2, 1e-2
-    ), "Refit Result is not correct. Refit failed"
+# # Check the output
+# # The refit happens while you call the mutable module again.
+# expected_outputs, refitted_outputs = model2(*inputs), mutable_module(*inputs)
+# for expected_output, refitted_output in zip(expected_outputs, refitted_outputs):
+#     assert torch.allclose(
+#         expected_output, refitted_output, 1e-2, 1e-2
+#     ), "Refit Result is not correct. Refit failed"
 
-print("Refit successfully!")
+# print("Refit successfully!")
 
-# %%
-# Saving Mutable Torch TensorRT Module
-# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# # %%
+# # Saving Mutable Torch TensorRT Module
+# # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-# Currently, saving is only enabled for C++ runtime, not python runtime.
-torch_trt.MutableTorchTensorRTModule.save(mutable_module, "mutable_module.pkl")
-reload = torch_trt.MutableTorchTensorRTModule.load("mutable_module.pkl")
+# # Currently, saving is only when "use_python" = False in settings
+# torch_trt.MutableTorchTensorRTModule.save(mutable_module, "mutable_module.pkl")
+# reload = torch_trt.MutableTorchTensorRTModule.load("mutable_module.pkl")
 
-# %%
-# Stable Diffusion with Huggingface
-# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# # %%
+# # Stable Diffusion with Huggingface
+# # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-# The LoRA checkpoint is from https://civitai.com/models/12597/moxin
+# # The LoRA checkpoint is from https://civitai.com/models/12597/moxin
 
-from diffusers import DiffusionPipeline
+# from diffusers import DiffusionPipeline
 
-with torch.no_grad():
-    settings = {
-        "use_python_runtime": True,
-        "enabled_precisions": {torch.float16},
-        "debug": True,
-        "immutable_weights": False,
-    }
+# with torch.no_grad():
+#     settings = {
+#         "use_python_runtime": True,
+#         "enabled_precisions": {torch.float16},
+#         "debug": True,
+#         "immutable_weights": False,
+#     }
 
-    model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-    device = "cuda:0"
+#     model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+#     device = "cuda:0"
 
-    prompt = "cinematic photo elsa, police uniform <lora:princess_xl_v2:0.8>, . 35mm photograph, film, bokeh, professional, 4k, highly detailed"
-    negative = "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, nude"
+#     prompt = "cinematic photo elsa, police uniform <lora:princess_xl_v2:0.8>, . 35mm photograph, film, bokeh, professional, 4k, highly detailed"
+#     negative = "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, nude"
 
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe.to(device)
+#     pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+#     pipe.to(device)
 
-    # The only extra line you need
-    pipe.unet = torch_trt.MutableTorchTensorRTModule(pipe.unet, **settings)
+#     # The only extra line you need
+#     pipe.unet = torch_trt.MutableTorchTensorRTModule(pipe.unet, **settings)
 
-    image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
-    image.save("./without_LoRA_mutable.jpg")
+#     image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
+#     image.save("./without_LoRA_mutable.jpg")
 
-    # Standard Huggingface LoRA loading procedure
-    pipe.load_lora_weights(
-        "stablediffusionapi/load_lora_embeddings",
-        weight_name="all-disney-princess-xl-lo.safetensors",
-        adapter_name="lora1",
-    )
-    pipe.set_adapters(["lora1"], adapter_weights=[1])
-    pipe.fuse_lora()
-    pipe.unload_lora_weights()
+#     # Standard Huggingface LoRA loading procedure
+#     pipe.load_lora_weights(
+#         "stablediffusionapi/load_lora_embeddings",
+#         weight_name="all-disney-princess-xl-lo.safetensors",
+#         adapter_name="lora1",
+#     )
+#     pipe.set_adapters(["lora1"], adapter_weights=[1])
+#     pipe.fuse_lora()
+#     pipe.unload_lora_weights()
 
-    # Refit triggered
-    image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
-    image.save("./with_LoRA_mutable.jpg")
+#     # Refit triggered
+#     image = pipe(prompt, negative_prompt=negative, num_inference_steps=30).images[0]
+#     image.save("./with_LoRA_mutable.jpg")
+
+
+# %%
+# Use Mutable Torch TensorRT module with dynamic shape
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+class Model(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, a, b, c={}):
+        x = torch.matmul(a, b)
+        x = torch.matmul(c["a"], c["b"].T)
+        x = 2 * c["b"][1]
+        return x
+
+
+model = Model().eval().cuda()
+inputs = (torch.rand(10, 3), torch.rand(3, 30))
+kwargs = {
+    "c": {"a": torch.rand(10, 30), "b": torch.rand(10, 30)},
+}
+
+dim = torch.export.Dim("dim", min=1, max=50)
+dim2 = torch.export.Dim("dim2", min=1, max=50)
+args_dynamic_shapes = ({1: dim}, {0: dim})
+kwarg_dynamic_shapes = {
+    "c": {"a": {}, "b": {0: dim2}},
+}
+# Export the model first with custom dynamic shape constraints
+# exp_program = torch.export.export(model, tuple(inputs), kwargs=k
+trt_gm = torch_trt.MutableTorchTensorRTModule(model, debug=True)
+trt_gm.set_expected_dynamic_shape_range(args_dynamic_shapes, kwarg_dynamic_shapes)
+# Run inference
+trt_gm(*inputs, **kwargs)