.. 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_advanced_usage.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_advanced_usage.py>`
        to download the full example code

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

.. _sphx_glr_tutorials__rendered_examples_dynamo_torch_compile_advanced_usage.py:


.. _torch_compile_advanced_usage:

Torch Compile Advanced Usage
======================================================

This interactive script is intended as an overview of the process by which `torch_tensorrt.compile(..., ir="torch_compile", ...)` works, and how it integrates with the `torch.compile` API.

.. GENERATED FROM PYTHON SOURCE LINES 10-12

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

.. GENERATED FROM PYTHON SOURCE LINES 12-16

.. code-block:: python


    import torch
    import torch_tensorrt


.. GENERATED FROM PYTHON SOURCE LINES 17-32

.. code-block:: python



    # We begin by defining a model
    class Model(torch.nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.relu = torch.nn.ReLU()

        def forward(self, x: torch.Tensor, y: torch.Tensor):
            x_out = self.relu(x)
            y_out = self.relu(y)
            x_y_out = x_out + y_out
            return torch.mean(x_y_out)



.. GENERATED FROM PYTHON SOURCE LINES 33-35

Compilation with `torch.compile` Using Default Settings
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 35-40

.. code-block:: python


    # Define sample float inputs and initialize model
    sample_inputs = [torch.rand((5, 7)).cuda(), torch.rand((5, 7)).cuda()]
    model = Model().eval().cuda()


.. GENERATED FROM PYTHON SOURCE LINES 41-49

.. code-block:: python


    # Next, we compile the model using torch.compile
    # For the default settings, we can simply call torch.compile
    # with the backend "torch_tensorrt", and run the model on an
    # input to cause compilation, as so:
    optimized_model = torch.compile(model, backend="torch_tensorrt")
    optimized_model(*sample_inputs)


.. GENERATED FROM PYTHON SOURCE LINES 50-52

Compilation with `torch.compile` Using Custom Settings
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 52-64

.. code-block:: python


    # First, we use Torch utilities to clean up the workspace
    # after the previous compile invocation
    torch._dynamo.reset()

    # Define sample half inputs and initialize model
    sample_inputs_half = [
        torch.rand((5, 7)).half().cuda(),
        torch.rand((5, 7)).half().cuda(),
    ]
    model_half = Model().eval().cuda()


.. GENERATED FROM PYTHON SOURCE LINES 65-88

.. code-block:: python


    # If we want to customize certain options in the backend,
    # but still use the torch.compile call directly, we can provide
    # custom options to the backend via the "options" keyword
    # which takes in a dictionary mapping options to values.
    #
    # For accepted backend options, see the CompilationSettings dataclass:
    # py/torch_tensorrt/dynamo/_settings.py
    backend_kwargs = {
        "enabled_precisions": {torch.half},
        "debug": True,
        "min_block_size": 2,
        "torch_executed_ops": {"torch.ops.aten.sub.Tensor"},
        "optimization_level": 4,
        "use_python_runtime": False,
    }

    # Run the model on an input to cause compilation, as so:
    optimized_model_custom = torch.compile(
        model_half, backend="torch_tensorrt", options=backend_kwargs
    )
    optimized_model_custom(*sample_inputs_half)


.. GENERATED FROM PYTHON SOURCE LINES 89-91

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

.. GENERATED FROM PYTHON SOURCE LINES 91-95

.. code-block:: python


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


.. GENERATED FROM PYTHON SOURCE LINES 96-105

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_advanced_usage.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_advanced_usage.py <torch_compile_advanced_usage.py>`

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

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


.. only:: html

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

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