Skip to content

Commit cfcd65d

Browse files
committed
chore: update dynamo export doc
Signed-off-by: Dheeraj Peri <[email protected]>
1 parent 83176fe commit cfcd65d

File tree

2 files changed

+74
-0
lines changed

2 files changed

+74
-0
lines changed

docsrc/index.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ User Guide
4040
------------
4141
* :ref:`creating_a_ts_mod`
4242
* :ref:`getting_started_with_fx`
43+
* :ref:`dynamo_export`
4344
* :ref:`ptq`
4445
* :ref:`runtime`
4546
* :ref:`saving_models`
@@ -54,6 +55,7 @@ User Guide
5455

5556
user_guide/creating_torchscript_module_in_python
5657
user_guide/getting_started_with_fx_path
58+
user_guide/dynamo_export
5759
user_guide/ptq
5860
user_guide/runtime
5961
user_guide/saving_models

docsrc/user_guide/dynamo_export.rst

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
.. _getting_started_with_dynamo:
2+
3+
Torch-TensorRT (Dynamo) Backend
4+
========================
5+
This guide presents Torch-TensorRT dynamo backend which compiles Pytorch programs
6+
into TensorRT engines through torch dynamo. Pytorch 2.1 introduced ``torch.export`` APIs which
7+
can export graphs from Pytorch programs using torch dynamo. Torch-TensorRT dynamo
8+
backend compiles these exported graphs and optimizes them using TensorRT. Here's a simple
9+
usage of the dynamo backend
10+
11+
.. code-block:: python
12+
13+
import torch
14+
import torch_tensorrt
15+
16+
model = MyModel().eval().cuda()
17+
inputs = torch.randn((1, 3, 224, 224), dtype=torch.float32).cuda()
18+
exp_program = torch.export(model, inputs)
19+
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs) # Output is a torch.fx.GraphModule
20+
trt_gm(inputs)
21+
22+
``torch_tensorrt.dynamo.compile`` is the main API for users to interact with Torch-TensorRT.
23+
The input type of the model should be ``ExportedProgram`` (ideally the output of torch.export) and output types is a ``torch.fx.GraphModule`` object.
24+
25+
Customizations
26+
---------------------------------------------
27+
28+
There are lot of options for users to customize their settings for optimizing with TensorRT.
29+
Some of the frequently used options are as follows:
30+
* inputs
31+
- For static shapes, this can be a list of a) torch tensors or b) `torch_tensorrt.Input` objects
32+
- For dynamic shapes, this should be a list of ``torch_tensorrt.Input`` objects.
33+
* enabled_precisions - Set of precisions that TensorRT builder can use during optimization.
34+
* truncate_long_and_double - Truncates long and double values to int and floats respectively.
35+
* torch_executed_ops - Operators which are forced to be executed by Torch.
36+
* min_block_size - Minimum number of consecutive operators required to be executed as a TensorRT segment.
37+
38+
The complete list of options can be found `here <https://github.com/pytorch/TensorRT/blob/123a486d6644a5bbeeec33e2f32257349acc0b8f/py/torch_tensorrt/dynamo/compile.py#L51-L77>`_
39+
Note: We do not support INT precision currently in Dynamo. Support for this currently exists in
40+
our Torchscript IR. We plan to implement similar support for dynamo in our next release.
41+
42+
Under the hood
43+
--------------
44+
45+
Under the hood, ``torch_tensorrt.dynamo.compile`` performs the following on the graph.
46+
47+
* Lowering - Applies lowering passes to add/remove operators for optimal conversion.
48+
* Partitioning - Partitions the graph into Pytorch and TensorRT segments based on the ``min_block_size`` and ``torch_executed_ops`` field.
49+
* Conversion - Pytorch ops get converted into TensorRT ops in this phase.
50+
* Optimization - Post conversion, we build the TensorRT engine and embed this inside the pytorch graph.
51+
52+
Tracing
53+
-------
54+
55+
``torch_tensorrt.dynamo.trace`` can be used to trace a Pytorch graphs and produce ``ExportedProgram``.
56+
This internally performs some decompositions of operators for downstream optimization.
57+
The ``ExportedProgram`` can then be used with ``torch_tensorrt.dynamo.compile`` API.
58+
If you have dynamic input shapes in your model, you can use this ``torch_tensorrt.dynamo.trace`` to export
59+
the model with dynamic shapes. Alternatively, you can use ``torch.export`` `with constraints <https://pytorch.org/docs/stable/export.html#expressing-dynamism>`_ directly as well.
60+
61+
.. code-block:: python
62+
63+
import torch
64+
import torch_tensorrt
65+
66+
inputs = torch_tensorrt.Input(min_shape=(1, 3, 224, 224),
67+
opt_shape=(4, 3, 224, 224),
68+
max_shape=(8, 3, 224, 224),
69+
dtype=torch.float32)
70+
model = MyModel().eval()
71+
exp_program = torch_tensorrt.dynamo.trace(model, inputs)
72+

0 commit comments

Comments
 (0)