Table of Contents
Shortcuts

# torch_tensorrt

## Functions

torch_tensorrt.compile(module: Any, ir: str = 'default', inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → Union[Module, ScriptModule, GraphModule, Callable[[...], Any]]

  Compile a PyTorch module for NVIDIA GPUs using TensorRT

  Takes a existing PyTorch module and a set of settings to configure the compiler and using the path specified in `ir` lower and compile the module to TensorRT returning a PyTorch Module back

  Converts specifically the forward method of a Module

  Parameters

    module (Union(torch.nn.Module,torch.jit.ScriptModule) – Source module

  Keyword Arguments

    - inputs (List[Union(Input, torch.Tensor)]) –

      Required List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum to select device type.

          inputs=[
              torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
              torch_tensorrt.Input(
                  min_shape=(1, 224, 224, 3),
                  opt_shape=(1, 512, 512, 3),
                  max_shape=(1, 1024, 1024, 3),
                  dtype=torch.int32
                  format=torch.channel_last
              ), # Dynamic input shape for input #2
              torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
          ]

    - arg_inputs (Tuple[Any, ...]) – Same as inputs. Alias for better understanding with kwarg_inputs.

    - kwarg_inputs (dict[Any, ...]) – Optional, kwarg inputs to the module forward function.

    - enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – The set of datatypes that TensorRT can use when selecting kernels

    - ir (str) – The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)

    - **kwargs – Additional settings for the specific requested strategy (See submodules for more info)

  Returns

    Compiled Module, when run it will execute via TensorRT

  Return type

    torch.nn.Module

torch_tensorrt.convert_method_to_trt_engine(module: Any, method_name: str = 'forward', inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, ir: str = 'default', enabled_precisions: Optional[Set[Union[dtype, dtype]]] = None, **kwargs: Any) → bytes

  Convert a TorchScript module method to a serialized TensorRT engine

  Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings

  Parameters

    module (Union(torch.nn.Module,torch.jit.ScriptModule) – Source module

  Keyword Arguments

    - inputs (List[Union(Input, torch.Tensor)]) –

      Required List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum to select device type.

          input=[
              torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
              torch_tensorrt.Input(
                  min_shape=(1, 224, 224, 3),
                  opt_shape=(1, 512, 512, 3),
                  max_shape=(1, 1024, 1024, 3),
                  dtype=torch.int32
                  format=torch.channel_last
              ), # Dynamic input shape for input #2
              torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
          ]

    - arg_inputs (Tuple[Any, ...]) – Same as inputs. Alias for better understanding with kwarg_inputs.

    - kwarg_inputs (dict[Any, ...]) – Optional, kwarg inputs to the module forward function.

    - enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – The set of datatypes that TensorRT can use when selecting kernels

    - ir (str) – The requested strategy to compile. (Options: default - Let Torch-TensorRT decide, ts - TorchScript with scripting path)

    - **kwargs – Additional settings for the specific requested strategy (See submodules for more info)

  Returns

    Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs

  Return type

    bytes

torch_tensorrt.save(module: Any, file_path: str = '', *, output_format: str = 'exported_program', inputs: Optional[Sequence[Tensor]] = None, arg_inputs: Optional[Sequence[Tensor]] = None, kwarg_inputs: Optional[dict[str, Any]] = None, retrace: bool = False) → None

  Save the model to disk in the specified output format.

  Parameters

    - module (Optional(torch.jit.ScriptModule | torch.export.ExportedProgram | torch.fx.GraphModule | CudaGraphsTorchTensorRTModule)) – Compiled Torch-TensorRT module

    - inputs (torch.Tensor) – Torch input tensors

    - arg_inputs (Tuple[Any, ...]) – Same as inputs. Alias for better understanding with kwarg_inputs.

    - kwarg_inputs (dict[Any, ...]) – Optional, kwarg inputs to the module forward function.

    - output_format (str) – Format to save the model. Options include exported_program | torchscript.

    - retrace (bool) – When the module type is a fx.GraphModule, this option re-exports the graph using torch.export.export(strict=False) to save it. This flag is experimental for now.

torch_tensorrt.load(file_path: str = '') → Any

  Load either a Torchscript model or ExportedProgram.

  Loads a TorchScript or ExportedProgram file from disk. File type will be detect the type using try, except.

  Parameters

    file_path (str) – Path to file on the disk

  Raises

    ValueError – If there is no file or the file is not either a TorchScript file or ExportedProgram file

## Classes

class torch_tensorrt.MutableTorchTensorRTModule(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = True, debug: bool = False, num_avg_timing_iters: int = 1, workspace_size: int = 0, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, truncate_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any)

  Initialize a MutableTorchTensorRTModule to seamlessly manipulate it like a regular PyTorch module. All TensorRT compilation and refitting processes are handled automatically as you work with the module. Any changes to its attributes or loading a different state_dict will trigger refitting or recompilation, which will be managed during the next forward pass.

  The MutableTorchTensorRTModule takes a PyTorch module and a set of configuration settings for the compiler. Once compilation is complete, the module maintains the connection between the TensorRT graph module and the original PyTorch module. Any modifications made to the MutableTorchTensorRTModule will be reflected in both the TensorRT graph module and the original PyTorch module.

  __init__(pytorch_model: Module, *, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Set[Union[dtype, dtype]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, immutable_weights: bool = True, debug: bool = False, num_avg_timing_iters: int = 1, workspace_size: int = 0, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, truncate_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', **kwargs: Any) → None
    Parameters

      pytorch_model (torch.nn.module) – Source module that needs to be accelerated

    Keyword Arguments

      - device (Union(Device, torch.device, dict)) –

        Target device for TensorRT engines to run on

            device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)

      - disable_tf32 (bool) – Force FP32 layers to use traditional as FP32 format vs the default behavior of rounding the inputs to 10-bit mantissas before multiplying, but accumulates the sum using 23-bit mantissas

      - assume_dynamic_shape_support (bool) – Setting this to true enables the converters work for both dynamic and static shapes. Default: False

      - sparse_weights (bool) – Enable sparsity for convolution and fully connected layers.

      - enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – The set of datatypes that TensorRT can use when selecting kernels

      - immutable_weights (bool) – Build non-refittable engines. This is useful for some layers that are not refittable.

      - debug (bool) – Enable debuggable engine

      - capability (EngineCapability) – Restrict kernel selection to safe gpu kernels or safe dla kernels

      - num_avg_timing_iters (python:int) – Number of averaging timing iterations used to select kernels

      - workspace_size (python:int) – Maximum size of workspace given to TensorRT

      - dla_sram_size (python:int) – Fast software managed RAM used by DLA to communicate within a layer.

      - dla_local_dram_size (python:int) – Host RAM used by DLA to share intermediate tensor data across operations

      - dla_global_dram_size (python:int) – Host RAM used by DLA to store weights and metadata for execution

      - truncate_double (bool) – Truncate weights provided in double (float64) to float32

      - calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – Calibrator object which will provide data to the PTQ system for INT8 Calibration

      - require_full_compilation (bool) – Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch

      - min_block_size (python:int) – The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT

      - torch_executed_ops (Collection[Target]) – Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but `require_full_compilation` is True

      - torch_executed_modules (List[str]) – List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but `require_full_compilation` is True

      - pass_through_build_failures (bool) – Error out if there are issues during compilation (only applicable to torch.compile workflows)

      - max_aux_stream (Optional[python:int]) – Maximum streams in the engine

      - version_compatible (bool) – Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines)

      - optimization_level – (Optional[int]): Setting a higher optimization level allows TensorRT to spend longer engine building time searching for more optimization options. The resulting engine may have better performance compared to an engine built with a lower optimization level. The default optimization level is 3. Valid values include integers from 0 to the maximum optimization level, which is currently 5. Setting it to be greater than the maximum level results in identical behavior to the maximum level.

      - use_python_runtime – (bool): Return a graph using a pure Python runtime, reduces options for serialization

      - use_fast_partitioner – (bool): Use the adjacency based partitioning scheme instead of the global partitioner. Adjacency partitioning is faster but may not be optimal. Use the global paritioner (`False`) if looking for best performance

      - enable_experimental_decompositions (bool) – Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT.

      - dryrun (bool) – Toggle for “Dryrun” mode, running everything except conversion to TRT and logging outputs

      - hardware_compatible (bool) – Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer)

      - timing_cache_path (str) – Path to the timing cache if it exists (or) where it will be saved after compilation

      - lazy_engine_init (bool) – Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime.

      - **kwargs – Any,

    Returns

      MutableTorchTensorRTModule

  compile() → None

    (Re)compile the TRT graph module using the PyTorch module. This function should be called whenever the weight structure get changed (shape, more layers…) MutableTorchTensorRTModule automatically catches weight value updates and call this function to recompile. If it fails to catch the changes, please call this function manually to recompile the TRT graph module.

  refit_gm() → None

    Refit the TRT graph module with any updates. This function should be called whenever the weight values get changed but the weight structure remains the same. MutableTorchTensorRTModule automatically catches weight value updates and call this function to refit the module. If it fails to catch the changes, please call this function manually to update the TRT graph module.

class torch_tensorrt.Input(*args: Any, **kwargs: Any)

  Defines an input to a module in terms of expected shape, data type and tensor format.

  Variables

    - shape_mode (torch_tensorrt.Input._ShapeMode) – Is input statically or dynamically shaped

    - shape (Tuple or Dict) –

      Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form

          {"min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple}

    - dtype (torch_tensorrt.dpython:type) – The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)

    - format (torch_tensorrt.TensorFormat) – The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)

  __init__(*args: Any, **kwargs: Any) → None

    __init__ Method for torch_tensorrt.Input

    Input accepts one of a few construction patterns

    Parameters

      shape (Tuple or List, optional) – Static shape of input tensor

    Keyword Arguments

      - shape (Tuple or List, optional) – Static shape of input tensor

      - min_shape (Tuple or List, optional) – Min size of input tensor’s shape range Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input’s shape_mode to DYNAMIC

      - opt_shape (Tuple or List, optional) – Opt size of input tensor’s shape range Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input’s shape_mode to DYNAMIC

      - max_shape (Tuple or List, optional) – Max size of input tensor’s shape range Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input’s shape_mode to DYNAMIC

      - dtype (torch.dpython:type or torch_tensorrt.dpython:type) – Expected data type for input tensor (default: torch_tensorrt.dtype.float32)

      - format (torch.memory_format or torch_tensorrt.TensorFormat) – The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)

      - tensor_domain (Tuple(python:float, python:float), optional) – The domain of allowed values for the tensor, as interval notation: [tensor_domain[0], tensor_domain[1]). Note: Entering “None” (or not specifying) will set the bound to [0, 2)

      - torch_tensor (torch.Tensor) – Holds a corresponding torch tensor with this Input.

      - name (str, optional) – Name of this input in the input nn.Module’s forward function. Used to specify dynamic shapes for the corresponding input in dynamo tracer.

    Examples

    - Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)

    - Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)

    - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW

  example_tensor(optimization_profile_field: Optional[str] = None) → Tensor

    Get an example tensor of the shape specified by the Input object

    Parameters

      optimization_profile_field (Optional(str)) – Name of the field to use for shape in the case the Input is dynamically shaped

    Returns

      A PyTorch Tensor

  classmethod from_tensor(t: Tensor, disable_memory_format_check: bool = False) → Input

    Produce a Input which contains the information of the given PyTorch tensor.

    Parameters

      - tensor (torch.Tensor) – A PyTorch tensor.

      - disable_memory_format_check (bool) – Whether to validate the memory formats of input tensors

    Returns

      A Input object.

  classmethod from_tensors(ts: Sequence[Tensor], disable_memory_format_check: bool = False) → List[Input]

    Produce a list of Inputs which contain the information of all the given PyTorch tensors.

    Parameters

      - tensors (Iterable[torch.Tensor]) – A list of PyTorch tensors.

      - disable_memory_format_check (bool) – Whether to validate the memory formats of input tensors

    Returns

      A list of Inputs.

  dtype: dtype = 1

    torch_tensorrt.dtype.float32)

    Type

      The expected data type of the input tensor (default

  format: memory_format = 1

    torch_tensorrt.memory_format.linear)

    Type

      The expected format of the input tensor (default

class torch_tensorrt.Device(*args: Any, **kwargs: Any)

  Defines a device that can be used to specify target devices for engines

  Variables

    - device_type (DeviceType) – Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.

    - gpu_id (python:int) – Device ID for target GPU

    - dla_core (python:int) – Core ID for target DLA core

    - allow_gpu_fallback (bool) – Whether falling back to GPU if DLA cannot support an op should be allowed

  __init__(*args: Any, **kwargs: Any)

    __init__ Method for torch_tensorrt.Device

    Device accepts one of a few construction patterns

    Parameters

      spec (str) – String with device spec e.g. “dla:0” for dla, core_id 0

    Keyword Arguments

      - gpu_id (python:int) – ID of target GPU (will get overridden if dla_core is specified to the GPU managing DLA). If specified, no positional arguments should be provided

      - dla_core (python:int) – ID of target DLA core. If specified, no positional arguments should be provided.

      - allow_gpu_fallback (bool) – Allow TensorRT to schedule operations on GPU if they are not supported on DLA (ignored if device type is not DLA)

    Examples

    - Device(“gpu:1”)

    - Device(“cuda:1”)

    - Device(“dla:0”, allow_gpu_fallback=True)

    - Device(gpu_id=0, dla_core=0, allow_gpu_fallback=True)

    - Device(dla_core=0, allow_gpu_fallback=True)

    - Device(gpu_id=1)

  device_type: DeviceType = 1

    Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.

  dla_core: int = -1

    Core ID for target DLA core

  gpu_id: int = -1

    Device ID for target GPU

## Enums

class torch_tensorrt.dtype(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

  Enum to describe data types to Torch-TensorRT, has compatibility with torch, tensorrt and numpy dtypes

  to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool = False) → Union[dtype, DataType, dtype, dtype]

    Convert dtype into the equivalent type in [torch, numpy, tensorrt]

    Converts `self` into one of numpy, torch, and tensorrt equivalent dtypes. If `self` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.dtype.try_to()`

    Parameters

      - t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – Data type enum from another library to convert to

      - use_default (bool) – In some cases a catch all type (such as `torch.float`) is sufficient, so instead of throwing an exception, return default value.

    Returns

      dtype equivalent `torch_tensorrt.dtype` from library enum `t`

    Return type

      Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)

    Raises

      TypeError – Unsupported data type or unknown target

    Examples

        # Succeeds
        float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

        # Failure
        float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Throws exception

  classmethod try_from(t: Union[dtype, DataType, dtype, dtype], use_default: bool = False) → Optional[dtype]

    Create a Torch-TensorRT dtype from another library’s dtype system.

    Takes a dtype enum from one of numpy, torch, and tensorrt and create a `torch_tensorrt.dtype`. If the source dtype system is not supported or the type is not supported in Torch-TensorRT, then returns `None`.

    Parameters

      - t (Union(torch.dpython:type, tensorrt.DataType, numpy.dpython:type, dpython:type)) – Data type enum from another library

      - use_default (bool) – In some cases a catch all type (such as `torch_tensorrt.dtype.f32`) is sufficient, so instead of throwing an exception, return default value.

    Returns

      Equivalent `torch_tensorrt.dtype` to `t` or `None`

    Return type

      Optional(dtype)

    Examples

        # Succeeds
        float_dtype = torch_tensorrt.dtype.try_from(torch.float) # Returns torch_tensorrt.dtype.f32

        # Unsupported type
        float_dtype = torch_tensorrt.dtype.try_from(torch.complex128) # Returns None

  try_to(t: Union[Type[dtype], Type[DataType], Type[dtype], Type[dtype]], use_default: bool) → Optional[Union[dtype, DataType, dtype, dtype]]

    Convert dtype into the equivalent type in [torch, numpy, tensorrt]

    Converts `self` into one of numpy, torch, and tensorrt equivalent dtypes. If `self` is not supported in the target library, then returns `None`.

    Parameters

      - t (Union(Type(torch.dpython:type), Type(tensorrt.DataType), Type(numpy.dpython:type), Type(dpython:type))) – Data type enum from another library to convert to

      - use_default (bool) – In some cases a catch all type (such as `torch.float`) is sufficient, so instead of throwing an exception, return default value.

    Returns

      dtype equivalent `torch_tensorrt.dtype` from library enum `t`

    Return type

      Optional(Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype))

    Examples

        # Succeeds
        float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float

        # Failure
        float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Returns None

  b

    Boolean value, equivalent to `dtype.bool`

  bf16

    16 bit “Brain” floating-point number, equivalent to `dtype.bfloat16`

  f16

    16 bit floating-point number, equivalent to `dtype.half`, `dtype.fp16` and `dtype.float16`

  f32

    32 bit floating-point number, equivalent to `dtype.float`, `dtype.fp32` and `dtype.float32`

  f64

    64 bit floating-point number, equivalent to `dtype.double`, `dtype.fp64` and `dtype.float64`

  f8

    8 bit floating-point number, equivalent to `dtype.fp8` and `dtype.float8`

  i32

    Signed 32 bit integer, equivalent to `dtype.int32` and `dtype.int`

  i64

    Signed 64 bit integer, equivalent to `dtype.int64` and `dtype.long`

  i8

    Signed 8 bit integer, equivalent to `dtype.int8`, when enabled as a kernel precision typically requires the model to support quantization

  u8

    Unsigned 8 bit integer, equivalent to `dtype.uint8`

  unknown

    Sentinel value

class torch_tensorrt.DeviceType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

  Type of device TensorRT will target

  to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Union[DeviceType, DeviceType]

    Convert `DeviceType` into the equivalent type in tensorrt

    Converts `self` into one of torch or tensorrt equivalent device type. If `self` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.DeviceType.try_to()`

    Parameters

      t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – Device type enum from another library to convert to

    Returns

      Device type equivalent `torch_tensorrt.DeviceType` in enum `t`

    Return type

      Union(tensorrt.DeviceType, DeviceType)

    Raises

      TypeError – Unknown target type or unsupported device type

    Examples

        # Succeeds
        trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA

  classmethod try_from(d: Union[DeviceType, DeviceType]) → Optional[DeviceType]

    Create a Torch-TensorRT device type enum from a TensorRT device type enum.

    Takes a device type enum from tensorrt and create a `torch_tensorrt.DeviceType`. If the source is not supported or the device type is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.DeviceType.try_from()`

    Parameters

      d (Union(tensorrt.DeviceType, DeviceType)) – Device type enum from another library

    Returns

      Equivalent `torch_tensorrt.DeviceType` to `d`

    Return type

      DeviceType

    Examples

        torchtrt_dla = torch_tensorrt.DeviceType._from(tensorrt.DeviceType.DLA)

  try_to(t: Union[Type[DeviceType], Type[DeviceType]], use_default: bool = False) → Optional[Union[DeviceType, DeviceType]]

    Convert `DeviceType` into the equivalent type in tensorrt

    Converts `self` into one of torch or tensorrt equivalent memory format. If `self` is not supported in the target library, then `None` will be returned.

    Parameters

      t (Union(Type(tensorrt.DeviceType), Type(DeviceType))) – Device type enum from another library to convert to

    Returns

      Device type equivalent `torch_tensorrt.DeviceType` in enum `t`

    Return type

      Optional(Union(tensorrt.DeviceType, DeviceType))

    Examples

        # Succeeds
        trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA

  DLA

    Target is a DLA core

  GPU

    Target is a GPU

  UNKNOWN

    Sentinel value

class torch_tensorrt.EngineCapability(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

  EngineCapability determines the restrictions of a network during build time and what runtime it targets.

  to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Union[EngineCapability, EngineCapability]

    Convert `EngineCapability` into the equivalent type in tensorrt

    Converts `self` into one of torch or tensorrt equivalent engine capability. If `self` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.EngineCapability.try_to()`

    Parameters

      t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – Engine capability enum from another library to convert to

    Returns

      Engine capability equivalent `torch_tensorrt.EngineCapability` in enum `t`

    Return type

      Union(tensorrt.EngineCapability, EngineCapability)

    Raises

      TypeError – Unknown target type or unsupported engine capability

    Examples

        # Succeeds
        torchtrt_dla_ec = torch_tensorrt.EngineCapability.DLA_STANDALONE.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA

  classmethod try_from() → Optional[EngineCapability]

    Create a Torch-TensorRT engine capability enum from a TensorRT engine capability enum.

    Takes a device type enum from tensorrt and create a `torch_tensorrt.EngineCapability`. If the source is not supported or the engine capability level is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.EngineCapability.try_from()`

    Parameters

      c (Union(tensorrt.EngineCapability, EngineCapability)) – Engine capability enum from another library

    Returns

      Equivalent `torch_tensorrt.EngineCapability` to `c`

    Return type

      EngineCapability

    Examples

        torchtrt_safety_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAEFTY)

  try_to(t: Union[Type[EngineCapability], Type[EngineCapability]]) → Optional[Union[EngineCapability, EngineCapability]]

    Convert `EngineCapability` into the equivalent type in tensorrt

    Converts `self` into one of torch or tensorrt equivalent engine capability. If `self` is not supported in the target library, then `None` will be returned.

    Parameters

      t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))) – Engine capability enum from another library to convert to

    Returns

      Engine capability equivalent `torch_tensorrt.EngineCapability` in enum `t`

    Return type

      Optional(Union(tensorrt.EngineCapability, EngineCapability))

    Examples

        # Succeeds
        trt_dla_ec = torch_tensorrt.EngineCapability.DLA.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA_STANDALONE

  DLA_STANDALONE

    `EngineCapability.DLA_STANDALONE` provides a restricted subset of network operations that are DLA compatible and the resulting serialized engine can be executed using standalone DLA runtime APIs.

  SAFETY

    EngineCapability.SAFETY provides a restricted subset of network operations that are safety certified and the resulting serialized engine can be executed with TensorRT’s safe runtime APIs in the tensorrt.safe namespace.

  STANDARD

    EngineCapability.STANDARD does not provide any restrictions on functionality and the resulting serialized engine can be executed with TensorRT’s standard runtime APIs.

class torch_tensorrt.memory_format(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)
  to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Union[memory_format, TensorFormat, memory_format]

    Convert `memory_format` into the equivalent type in torch or tensorrt

    Converts `self` into one of torch or tensorrt equivalent memory format. If `self` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly.

    Alternatively use `torch_tensorrt.memory_format.try_to()`

    Parameters

      t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – Memory format type enum from another library to convert to

    Returns

      Memory format equivalent `torch_tensorrt.memory_format` in enum `t`

    Return type

      Union(torch.memory_format, tensorrt.TensorFormat, memory_format)

    Raises

      TypeError – Unknown target type or unsupported memory format

    Examples

        # Succeeds
        tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous

  classmethod try_from(f: Union[memory_format, TensorFormat, memory_format]) → Optional[memory_format]

    Create a Torch-TensorRT memory format enum from another library memory format enum.

    Takes a memory format enum from one of torch, and tensorrt and create a `torch_tensorrt.memory_format`. If the source is not supported or the memory format is not supported in Torch-TensorRT, then `None` will be returned.

    Parameters

      f (Union(torch.memory_format, tensorrt.TensorFormat, memory_format)) – Memory format enum from another library

    Returns

      Equivalent `torch_tensorrt.memory_format` to `f`

    Return type

      Optional(memory_format)

    Examples

        torchtrt_linear = torch_tensorrt.memory_format.try_from(torch.contiguous)

  try_to(t: Union[Type[memory_format], Type[TensorFormat], Type[memory_format]]) → Optional[Union[memory_format, TensorFormat, memory_format]]

    Convert `memory_format` into the equivalent type in torch or tensorrt

    Converts `self` into one of torch or tensorrt equivalent memory format. If `self` is not supported in the target library, then `None` will be returned

    Parameters

      t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))) – Memory format type enum from another library to convert to

    Returns

      Memory format equivalent `torch_tensorrt.memory_format` in enum `t`

    Return type

      Optional(Union(torch.memory_format, tensorrt.TensorFormat, memory_format))

    Examples

        # Succeeds
        tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous

  cdhw32

    Thirty-two wide channel vectorized row major format with 3 spatial dimensions.

    This format is bound to FP16 and INT8. It is only available for dimensions >= 4.

    For a tensor with dimensions {N, C, D, H, W}, the memory layout is equivalent to a C array with dimensions [N][(C+31)/32][D][H][W][32], with the tensor coordinates (n, d, c, h, w) mapping to array subscript [n][c/32][d][h][w][c%32].

  chw16

    Sixteen wide channel vectorized row major format.

    This format is bound to FP16. It is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to a C array with dimensions [N][(C+15)/16][H][W][16], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][c/16][h][w][c%16].

  chw2

    Two wide channel vectorized row major format.

    This format is bound to FP16 in TensorRT. It is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to a C array with dimensions [N][(C+1)/2][H][W][2], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][c/2][h][w][c%2].

  chw32

    Thirty-two wide channel vectorized row major format.

    This format is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to a C array with dimensions [N][(C+31)/32][H][W][32], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][c/32][h][w][c%32].

  chw4

    Four wide channel vectorized row major format. This format is bound to INT8. It is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to a C array with dimensions [N][(C+3)/4][H][W][4], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][c/4][h][w][c%4].

  dhwc

    Non-vectorized channel-last format. This format is bound to FP32. It is only available for dimensions >= 4.

    Equivient to `memory_format.channels_last_3d`

  dhwc8

    Eight channel format where C is padded to a multiple of 8.

    This format is bound to FP16, and it is only available for dimensions >= 4.

    For a tensor with dimensions {N, C, D, H, W}, the memory layout is equivalent to an array with dimensions [N][D][H][W][(C+7)/8*8], with the tensor coordinates (n, c, d, h, w) mapping to array subscript [n][d][h][w][c].

  dla_hwc4

    DLA image format. channel-last format. C can only be 1, 3, 4. If C == 3 it will be rounded to 4. The stride for stepping along the H axis is rounded up to 32 bytes.

    This format is bound to FP16/Int8 and is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, with C’ is 1, 4, 4 when C is 1, 3, 4 respectively, the memory layout is equivalent to a C array with dimensions [N][H][roundUp(W, 32/C’/elementSize)][C’] where elementSize is 2 for FP16 and 1 for Int8, C’ is the rounded C. The tensor coordinates (n, c, h, w) maps to array subscript [n][h][w][c].

  dla_linear

    DLA planar format. Row major format. The stride for stepping along the H axis is rounded up to 64 bytes.

    This format is bound to FP16/Int8 and is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to a C array with dimensions [N][C][H][roundUp(W, 64/elementSize)] where elementSize is 2 for FP16 and 1 for Int8, with the tensor coordinates (n, c, h, w) mapping to array subscript [n][c][h][w].

  hwc

    Non-vectorized channel-last format. This format is bound to FP32 and is only available for dimensions >= 3.

    Equivient to `memory_format.channels_last`

  hwc16

    Sixteen channel format where C is padded to a multiple of 16. This format is bound to FP16. It is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to the array with dimensions [N][H][W][(C+15)/16*16], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][h][w][c].

  hwc8

    Eight channel format where C is padded to a multiple of 8.

    This format is bound to FP16. It is only available for dimensions >= 3.

    For a tensor with dimensions {N, C, H, W}, the memory layout is equivalent to the array with dimensions [N][H][W][(C+7)/8*8], with the tensor coordinates (n, c, h, w) mapping to array subscript [n][h][w][c].

  linear

    Row major linear format.

    For a tensor with dimensions {N, C, H, W}, the W axis always has unit stride, and the stride of every other axis is at least the product of the next dimension times the next stride. the strides are the same as for a C array with dimensions [N][C][H][W].

    Equivient to `memory_format.contiguous`

## Submodules

- torch_tensorrt.logging
- torch_tensorrt.ts.ptq
- torch_tensorrt.ts
- torch_tensorrt.fx
- torch_tensorrt.dynamo
- torch_tensorrt.runtime

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.