|
3 | 3 | """
|
4 | 4 | (Beta) ``torch.export`` AOTInductor Tutorial for Python runtime
|
5 | 5 | ===================================================
|
6 |
| -**Author:** Ankith Gunapal, Bin Bao |
| 6 | +**Author:** Ankith Gunapal, Bin Bao, Angela Yi |
7 | 7 | """
|
8 | 8 |
|
9 | 9 | ######################################################################
|
10 | 10 | #
|
11 | 11 | # .. warning::
|
12 | 12 | #
|
13 |
| -# ``torch._export.aot_compile`` and ``torch._export.aot_load`` are in Beta status and are subject to backwards compatibility |
| 13 | +# ``torch._inductor.aot_compile`` and ``torch._export.aot_load`` are in Beta status and are subject to backwards compatibility |
14 | 14 | # breaking changes. This tutorial provides an example of how to use these APIs for model deployment using Python runtime.
|
15 | 15 | #
|
16 | 16 | # It has been shown `previously <https://pytorch.org/docs/stable/torch.compiler_aot_inductor.html#>`__ how AOTInductor can be used
|
|
19 | 19 | #
|
20 | 20 | #
|
21 | 21 | # In this tutorial, you will learn an end-to-end example of how to use AOTInductor for python runtime.
|
22 |
| -# We will look at how to use :func:`torch._export.aot_compile` to generate a shared library. |
23 |
| -# Additionally, we will examine how to execute the shared library in Python runtime using :func:`torch._export.aot_load`. |
| 22 | +# We will look at how to use :func:`torch._inductor.aot_compile` along with :func:`torch.export.export` to generate a |
| 23 | +# shared library. Additionally, we will examine how to execute the shared library in Python runtime using :func:`torch._export.aot_load`. |
24 | 24 | # You will learn about the speed up seen in the first inference time using AOTInductor, especially when using
|
25 | 25 | # ``max-autotune`` mode which can take some time to execute.
|
26 | 26 | #
|
|
33 | 33 | # Prerequisites
|
34 | 34 | # -------------
|
35 | 35 | # * PyTorch 2.4 or later
|
36 |
| -# * Basic understanding of ``torch._export`` and AOTInductor |
| 36 | +# * Basic understanding of ``torch.export`` and AOTInductor |
37 | 37 | # * Complete the `AOTInductor: Ahead-Of-Time Compilation for Torch.Export-ed Models <https://pytorch.org/docs/stable/torch.compiler_aot_inductor.html#>`_ tutorial
|
38 | 38 |
|
39 | 39 | ######################################################################
|
40 | 40 | # What you will learn
|
41 | 41 | # ----------------------
|
42 | 42 | # * How to use AOTInductor for python runtime.
|
43 |
| -# * How to use :func:`torch._export.aot_compile` to generate a shared library |
| 43 | +# * How to use :func:`torch._inductor.aot_compile` along with :func:`torch.export.export` to generate a shared library |
44 | 44 | # * How to run a shared library in Python runtime using :func:`torch._export.aot_load`.
|
45 | 45 | # * When do you use AOTInductor for python runtime
|
46 | 46 |
|
|
49 | 49 | # ------------
|
50 | 50 | #
|
51 | 51 | # We will use the TorchVision pretrained `ResNet18` model and TorchInductor on the
|
52 |
| -# exported PyTorch program using :func:`torch._export.aot_compile`. |
| 52 | +# exported PyTorch program using :func:`torch._inductor.aot_compile`. |
53 | 53 | #
|
54 | 54 | # .. note::
|
55 | 55 | #
|
|
115 | 115 | # .. note::
|
116 | 116 | #
|
117 | 117 | # In the example above, we specified ``batch_size=1`` for inference and it still functions correctly even though we specified ``min=2`` in
|
118 |
| -# :func:`torch._export.aot_compile`. |
| 118 | +# :func:`torch.export.export`. |
119 | 119 |
|
120 | 120 |
|
121 | 121 | import os
|
|
139 | 139 | # model deployment using Python.
|
140 | 140 | # There are mainly two reasons why you would use AOTInductor Python Runtime:
|
141 | 141 | #
|
142 |
| -# - ``torch._export.aot_compile`` generates a shared library. This is useful for model |
| 142 | +# - ``torch._inductor.aot_compile`` generates a shared library. This is useful for model |
143 | 143 | # versioning for deployments and tracking model performance over time.
|
144 | 144 | # - With :func:`torch.compile` being a JIT compiler, there is a warmup
|
145 | 145 | # cost associated with the first compilation. Your deployment needs to account for the
|
146 | 146 | # compilation time taken for the first inference. With AOTInductor, the compilation is
|
147 |
| -# done offline using ``torch._export.aot_compile``. The deployment would only load the |
148 |
| -# shared library using ``torch._export.aot_load`` and run inference. |
| 147 | +# done offline using ``torch.export.export`` & ``torch._indutor.aot_compile``. The deployment |
| 148 | +# would only load the shared library using ``torch._export.aot_load`` and run inference. |
149 | 149 | #
|
150 | 150 | #
|
151 | 151 | # The section below shows the speedup achieved with AOTInductor for first inference
|
@@ -218,7 +218,7 @@ def timed(fn):
|
218 | 218 | # ----------
|
219 | 219 | #
|
220 | 220 | # In this recipe, we have learned how to effectively use the AOTInductor for Python runtime by
|
221 |
| -# compiling and loading a pretrained ``ResNet18`` model using the ``torch._export.aot_compile`` |
| 221 | +# compiling and loading a pretrained ``ResNet18`` model using the ``torch._inductor.aot_compile`` |
222 | 222 | # and ``torch._export.aot_load`` APIs. This process demonstrates the practical application of
|
223 | 223 | # generating a shared library and running it within a Python environment, even with dynamic shape
|
224 | 224 | # considerations and device-specific optimizations. We also looked at the advantage of using
|
|
0 commit comments