.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/_rendered_examples/dynamo/torch_compile_resnet_example.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_tutorials__rendered_examples_dynamo_torch_compile_resnet_example.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_tutorials__rendered_examples_dynamo_torch_compile_resnet_example.py:


.. _torch_compile_resnet:

Compiling ResNet using the Torch-TensorRT `torch.compile` Backend
==========================================================

This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a ResNet model.

.. GENERATED FROM PYTHON SOURCE LINES 10-12

Imports and Model Definition
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 12-17

.. code-block:: python


    import torch
    import torch_tensorrt
    import torchvision.models as models


.. GENERATED FROM PYTHON SOURCE LINES 18-23

.. code-block:: python


    # Initialize model with half precision and sample inputs
    model = models.resnet18(pretrained=True).half().eval().to("cuda")
    inputs = [torch.randn((1, 3, 224, 224)).to("cuda").half()]


.. GENERATED FROM PYTHON SOURCE LINES 24-26

Optional Input Arguments to `torch_tensorrt.compile`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 26-43

.. code-block:: python


    # Enabled precision for TensorRT optimization
    enabled_precisions = {torch.half}

    # Whether to print verbose logs
    debug = True

    # Workspace size for TensorRT
    workspace_size = 20 << 30

    # Maximum number of TRT Engines
    # (Lower value allows more graph segmentation)
    min_block_size = 7

    # Operations to Run in Torch, regardless of converter support
    torch_executed_ops = {}


.. GENERATED FROM PYTHON SOURCE LINES 44-46

Compilation with `torch_tensorrt.compile`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 46-59

.. code-block:: python


    # Build and compile the model with torch.compile, using Torch-TensorRT backend
    optimized_model = torch_tensorrt.compile(
        model,
        ir="torch_compile",
        inputs=inputs,
        enabled_precisions=enabled_precisions,
        debug=debug,
        workspace_size=workspace_size,
        min_block_size=min_block_size,
        torch_executed_ops=torch_executed_ops,
    )


.. GENERATED FROM PYTHON SOURCE LINES 60-62

Equivalently, we could have run the above via the torch.compile frontend, as so:
`optimized_model = torch.compile(model, backend="torch_tensorrt", options={"enabled_precisions": enabled_precisions, ...}); optimized_model(*inputs)`

.. GENERATED FROM PYTHON SOURCE LINES 64-66

Inference
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 66-71

.. code-block:: python


    # Does not cause recompilation (same batch size as input)
    new_inputs = [torch.randn((1, 3, 224, 224)).half().to("cuda")]
    new_outputs = optimized_model(*new_inputs)


.. GENERATED FROM PYTHON SOURCE LINES 72-77

.. code-block:: python


    # Does cause recompilation (new batch size)
    new_batch_size_inputs = [torch.randn((8, 3, 224, 224)).half().to("cuda")]
    new_batch_size_outputs = optimized_model(*new_batch_size_inputs)


.. GENERATED FROM PYTHON SOURCE LINES 78-80

Cleanup
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 80-84

.. code-block:: python


    # Finally, we use Torch utilities to clean up the workspace
    torch._dynamo.reset()


.. GENERATED FROM PYTHON SOURCE LINES 85-94

Cuda Driver Error Note
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Occasionally, upon exiting the Python runtime after Dynamo compilation with `torch_tensorrt`,
one may encounter a Cuda Driver Error. This issue is related to https://github.com/NVIDIA/TensorRT/issues/2052
and can be resolved by wrapping the compilation/inference in a function and using a scoped call, as in::

      if __name__ == '__main__':
          compile_engine_and_infer()


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_tutorials__rendered_examples_dynamo_torch_compile_resnet_example.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example




    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: torch_compile_resnet_example.py <torch_compile_resnet_example.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: torch_compile_resnet_example.ipynb <torch_compile_resnet_example.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_