DOCSEARCHbyAlgolia
Submit Docs

torch2trt

repository·master·Indexed 26 days ago

https://github.com/nvidia-ai-iot/torch2trt

A PyTorch to TensorRT converter that utilizes the TensorRT Python API. It allows for easy conversion of PyTorch modules to TensorRT engines via the torch2trt function and supports extensibility through Python-based custom layer converters using the @tensorrt_converter decorator. The library includes optional plugins for unsupported layers, an ONNX-based conversion workflow, and experimental community features such as Quantization Aware Training (QAT) for INT8 conversion of Conv2d and ReLU layers.

Tokens
7.2K
Snippets
16
Records
25
Agent score
85%

What's inside torch2trt

  1. Overview of torch2trt

    master

    torch2trt is a tool designed to convert PyTorch modules to TensorRT engines using the TensorRT Python API. It provides two primary benefits for developers:

    1. Ease of Use: You can convert PyTorch modules using a single function call: torch2trt.
    2. Extensibility: You can implement custom layer converters in Python and register them using the @tensorrt_converter decorator.

  2. Understand Quantization Aware Training (QAT) layer implementations

    master

    In torch2trt, every QAT layer provides two distinct implementations:

    1. Training Implementation: Performs quantization of weights and activations during the forward pass. This is used during the training phase to simulate quantization effects.
    2. Inference Implementation: Designed specifically for TensorRT conversion. Instead of performing quantization operations in the forward pass (which would result in unwanted operations in the final TensorRT engine), the inference version carries only the learned parameters, such as zero point and scale, to be used during the conversion process.

    This dual-implementation approach ensures that the final TensorRT engine remains efficient and contains only the necessary optimized layers.

  3. Save and load a converted TRTModule

    master

    You can persist a converted TensorRT model by saving its state_dict. To reload it, instantiate a TRTModule and use load_state_dict with the loaded weights.

    # Save the model as a state_dict
torch.save(model_trt.state_dict(), 'alexnet_trt.pth')

# Load the saved model into a TRTModule
from torch2trt import TRTModule

model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))

  4. Enable FP16 Precision in torch2trt

    master

    To improve throughput using FP16 precision, set the fp16_mode parameter to True during conversion. This allows the TensorRT optimizer to automatically select layers with FP16 precision for better performance. Note that setting this to True does not guarantee all layers will be FP16; the optimizer selects the best tactics for performance.

    To ensure the input and output bindings are also FP16, you must convert your PyTorch model and input data to half precision using .half() before calling torch2trt.

    # For FP32 input/output bindings with internal FP16 optimization
model = model.float()
data = data.float()
model_trt = torch2trt(model, [data], fp16_mode=True)

# For FP16 input/output bindings and internal FP16 optimization
model = model.half()
data = data.half()
model_trt = torch2trt(model, [data], fp16_mode=True)

  5. Install torch2trt without plugins

    master

    To install the torch2trt library without compiling additional plugins, clone the repository and run the setup script using Python.

    Prerequisites:

    • torch2trt depends on the TensorRT Python API.
    • On Jetson, this is included in the latest JetPack.
    • On Desktop, you must follow the TensorRT Installation Guide or use an NVIDIA NGC PyTorch docker container.
    git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

  6. Enable INT8 Precision in torch2trt

    master

    To enable INT8 precision, set int8_mode=True. Because INT8 can significantly impact accuracy, calibration is required.

    Calibration Methods

    1. Input Data Calibration: By default, torch2trt uses the inputs provided to the function for calibration. Use this for small datasets that fit in memory.

    2. Dataset Calibration: For larger datasets, use the int8_calib_dataset parameter. You must provide a class that implements __len__ (returning the number of samples) and __getitem__ (returning a list of input tensors matching the model's input shapes).

    Calibration Configuration

    • Algorithm: Override the default algorithm using int8_calib_algorithm with a tensorrt.CalibrationAlgoType value.
    • Batch Size: Control the number of samples pulled during calibration using int8_calib_batch_size.
    # Method 1: Calibrate using provided input data
data = torch.randn(64, 3, 224, 224).cuda().eval()
model_trt = torch2trt(model, [data], int8_mode=True)

# Method 2: Calibrate using a custom dataset object
class ImageFolderCalibDataset():
    def __init__(self, root):
        self.dataset = ImageFolder(root=root, transform=Compose([...]))
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        image, _ = self.dataset[idx]
        image = image[None, ...]
        return [image]

dataset = ImageFolderCalibDataset('images')
model_trt = torch2trt(model, [data], int8_calib_dataset=dataset)

# Method 3: Custom algorithm and batch size
import tensorrt as trt
model_trt = torch2trt(model, [data], int8_mode=True, int8_calib_algorithm=trt.CalibrationAlgoType.MINMAX_CALIBRATION, int8_calib_batch_size=32)

  7. Install experimental community features (optional)

    master

    To install experimental features under torch2trt.contrib, such as Quantization Aware Training (QAT) (requires TensorRT>=7.0), run the build script located in the scripts directory.

    git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts    
bash build_contrib.sh

  8. Convert a PyTorch model to TensorRT using torch2trt

    master

    To convert a PyTorch module to TensorRT, use the torch2trt function. You must provide the PyTorch model and a list containing example input data (tensors) that match the expected input shapes.

    Note: Once converted, you must use the same input shapes during execution. The only exception is the batch size, which can vary up to the value specified by the max_batch_size parameter.

    import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()

# create example data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

  9. Install the torch2trt plugins library (optional)

    master

    To add support for layers not natively supported by TensorRT, install the torch2trt plugins library using CMake. Once installed and found by the system, the associated layer converters are implicitly enabled.

    Note: Plugins are now maintained as an independent library. If you require the deprecated plugins that depend on PyTorch, use python setup.py install --plugins instead.

    cmake -B build . && cmake --build build --target install && ldconfig

  10. Implement a custom converter in torch2trt

    master

    You can extend torch2trt by implementing custom converters for specific PyTorch functional calls. This is done by defining a function and decorating it with @tensorrt_converter('path.to.pytorch.function').

    How it works

    torch2trt intercepts the specified PyTorch function call. It passes a ConversionContext object to your converter, which provides access to the TensorRT network being built and the arguments/return values of the original function.

    To link PyTorch tensors to TensorRT, use the ._trt attribute on input tensors and assign the resulting TensorRT output to the ._trt attribute of the output tensor.

    ConversionContext Attributes

    • ctx.network: The TensorRT network currently being constructed.
    • ctx.method_args: Positional arguments passed to the PyTorch function. Input tensors in this list will have a ._trt attribute.
    • ctx.method_kwargs: Keyword arguments passed to the PyTorch function.
    • ctx.method_return: The value returned by the PyTorch function. You must set the ._trt attribute on this object to represent the TensorRT output.
    import tensorrt as trt
from torch2trt import tensorrt_converter

@tensorrt_converter('torch.nn.ReLU.forward')
def convert_ReLU(ctx):
    input = ctx.method_args[1]
    output = ctx.method_return
    layer = ctx.network.add_activation(input=input._trt, type=trt.ActivationType.RELU)  
    output._trt = layer.get_output(0)

  11. Use Quantization Aware Training (QAT) for INT8 conversion

    master

    The contrib/qat module provides specialized layers and converters designed for Quantization Aware Training (QAT). This allows you to convert specific layers into INT8 format to optimize model performance for TensorRT deployment.

    Supported Layers

    • Conv2d
    • Conv2d with fused Batch Normalization (Conv2d + fused BN)
    • ReLU

    Supported Quantization Techniques

    • Per tensor quantization
    • Symmetric quantization

    Roadmap

    • Planned Layers: Pooling layers, Linear layers.
    • Planned Techniques: Per channel quantization, asymmetric quantization.

  12. Install the torch2trt Python library

    master

    To install the core torch2trt Python library, clone the repository and run the setup script. Note that torch2trt depends on the TensorRT Python API, which is included in JetPack for Jetson devices or must be installed manually for desktop environments.

    git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install