.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/_rendered_examples/dynamo/dynamo_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_dynamo_compile_advanced_usage.py>`
        to download the full example code

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

.. _sphx_glr_tutorials__rendered_examples_dynamo_dynamo_compile_advanced_usage.py:


.. _dynamo_compile_advanced_usage:

Dynamo Compile Advanced Usage
======================================================

This interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API.

.. GENERATED FROM PYTHON SOURCE LINES 10-12

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

.. GENERATED FROM PYTHON SOURCE LINES 12-17

.. code-block:: python


    import torch
    from torch_tensorrt.dynamo.backend import create_backend
    from torch_tensorrt.fx.lower_setting import LowerPrecision


.. GENERATED FROM PYTHON SOURCE LINES 18-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 "tensorrt", and run the model on an
    # input to cause compilation, as so:
    optimized_model = torch.compile(model, backend="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-60

.. code-block:: python


    # 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 61-77

.. code-block:: python


    # If we want to customize certain options in the backend,
    # but still use the torch.compile call directly, we can call the
    # convenience/helper function create_backend to create a custom backend
    # which has been pre-populated with certain keys
    custom_backend = create_backend(
        lower_precision=LowerPrecision.FP16,
        debug=True,
        min_block_size=2,
        torch_executed_ops={},
    )

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


.. GENERATED FROM PYTHON SOURCE LINES 78-80

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

.. GENERATED FROM PYTHON SOURCE LINES 80-86

.. code-block:: python


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

    with torch.no_grad():
        torch.cuda.empty_cache()


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

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


.. _sphx_glr_download_tutorials__rendered_examples_dynamo_dynamo_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: dynamo_compile_advanced_usage.py <dynamo_compile_advanced_usage.py>`

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

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


.. only:: html

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

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