|
| 1 | +""" |
| 2 | +.. _refit_engine_example: |
| 3 | +
|
| 4 | +Refit TenorRT Graph Module with Torch-TensorRT |
| 5 | +=================================================================== |
| 6 | +
|
| 7 | +We are going to demonstrate how a compiled TensorRT Graph Module can be refitted with updated weights. |
| 8 | +
|
| 9 | +In many cases, we frequently update the weights of models, such as applying various LoRA to Stable Diffusion or constant A/B testing of AI products. |
| 10 | +That poses challenges for TensorRT inference optimizations, as compiling the TensorRT engines takes significant time, making repetitive compilation highly inefficient. |
| 11 | +Torch-TensorRT supports refitting TensorRT graph modules without re-compiling the engine, considerably accelerating the workflow. |
| 12 | +
|
| 13 | +In this tutorial, we are going to walk through |
| 14 | +1. Compiling a PyTorch model to a TensorRT Graph Module |
| 15 | +2. Save and load a graph module |
| 16 | +3. Refit the graph module |
| 17 | +""" |
| 18 | + |
| 19 | +# %% |
| 20 | +# Standard Workflow |
| 21 | +# ----------------------------- |
| 22 | + |
| 23 | +# %% |
| 24 | +# Imports and model definition |
| 25 | +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| 26 | + |
| 27 | +import numpy as np |
| 28 | +import torch |
| 29 | +import torch_tensorrt as torch_trt |
| 30 | +import torchvision.models as models |
| 31 | +from torch_tensorrt.dynamo import refit_module_weights |
| 32 | + |
| 33 | +np.random.seed(0) |
| 34 | +torch.manual_seed(0) |
| 35 | +inputs = [torch.rand((1, 3, 224, 224)).to("cuda")] |
| 36 | + |
| 37 | + |
| 38 | +# %% |
| 39 | +# Compile the module for the first time and save it. |
| 40 | +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| 41 | + |
| 42 | +model = models.resnet18(pretrained=False).eval().to("cuda") |
| 43 | +exp_program = torch.export.export(model, tuple(inputs)) |
| 44 | +enabled_precisions = {torch.float} |
| 45 | +debug = False |
| 46 | +workspace_size = 20 << 30 |
| 47 | +min_block_size = 0 |
| 48 | +use_python_runtime = False |
| 49 | +torch_executed_ops = {} |
| 50 | +trt_gm = torch_trt.dynamo.compile( |
| 51 | + exp_program, |
| 52 | + tuple(inputs), |
| 53 | + use_python_runtime=use_python_runtime, |
| 54 | + enabled_precisions=enabled_precisions, |
| 55 | + debug=debug, |
| 56 | + min_block_size=min_block_size, |
| 57 | + torch_executed_ops=torch_executed_ops, |
| 58 | + make_refitable=True, |
| 59 | +) # Output is a torch.fx.GraphModule |
| 60 | + |
| 61 | +# Save the graph module as an exported program |
| 62 | +# This is only supported when use_python_runtime = False |
| 63 | +torch_trt.save(trt_gm, "./compiled.ep", inputs=inputs) |
| 64 | + |
| 65 | + |
| 66 | +# %% |
| 67 | +# Refit the module with update model weights |
| 68 | +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ |
| 69 | + |
| 70 | +# Create and compile the updated model |
| 71 | +model2 = models.resnet18(pretrained=True).eval().to("cuda") |
| 72 | +exp_program2 = torch.export.export(model2, tuple(inputs)) |
| 73 | + |
| 74 | + |
| 75 | +compiled_trt_ep = torch_trt.load("./compiled.ep") |
| 76 | + |
| 77 | +# This returns a new module with updated weights |
| 78 | +new_trt_gm = refit_module_weights( |
| 79 | + compiled_module=compiled_trt_ep, |
| 80 | + new_weight_module=exp_program2, |
| 81 | + inputs=inputs, |
| 82 | +) |
| 83 | + |
| 84 | +# Check the output |
| 85 | +expected_outputs, refitted_outputs = exp_program2.module()(*inputs), new_trt_gm(*inputs) |
| 86 | +for expected_output, refitted_output in zip(expected_outputs, refitted_outputs): |
| 87 | + assert torch.allclose( |
| 88 | + expected_output, refitted_output, 1e-2, 1e-2 |
| 89 | + ), "Refit Result is not correct. Refit failed" |
| 90 | + |
| 91 | +print("Refit successfully!") |
| 92 | + |
| 93 | +# %% |
| 94 | +# Alterative Workflow using Python Runtime |
| 95 | +# ----------------------------- |
| 96 | + |
| 97 | +# Currently python runtime does not support engine serialization. So the refitting will be done in the same runtime. |
| 98 | +# This usecase is more useful when you need to switch different weights in the same runtime, such as using Stable Diffusion. |
0 commit comments