Coggle数据科学 | Kaggle 知识点：Torch-TensorRT加速推理

本文来源公众号“Coggle数据科学”，仅用于学术分享，侵权删，干货满满。

原文链接：Kaggle 知识点：Torch-TensorRT加速推理

Torch-TensorRT是一个专为PyTorch设计的推理编译器，它通过NVIDIA的TensorRT深度学习优化器和运行时，为NVIDIA GPU提供优化的推理代码。这个工具的目的是提高在NVIDIA GPU上运行PyTorch模型的效率，同时保持开发过程的简便性。

安装Torch-TensorRT

https://pytorch.org/TensorRT/getting_started/installation.html

Torch-TensorRT 2.x主要围绕Python进行开发，因此可以在pypi.org上找到预编译的版本，而不需要从源代码编译。

python -m pip install torch torch-tensorrt tensorrt

基础使用

方法1：使用torch.compile

import torch
import torch_tensorrt

model = MyModel().eval().cuda() # define your model here
x = torch.randn((1, 3, 224, 224)).cuda() # define what the inputs to the model will look like

optimized_model = torch.compile(model, backend="tensorrt")
optimized_model(x) # compiled on first run

optimized_model(x) # this will be fast!

方法2：导出模型torch.export

import torch
import torch_tensorrt

model = MyModel().eval().cuda() # define your model here
inputs = [torch.randn((1, 3, 224, 224)).cuda()] # define a list of representative inputs here

trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs)
torch_tensorrt.save(trt_gm, "trt.ep", inputs=inputs) # PyTorch only supports Python runtime for an ExportedProgram. For C++ deployment, use a TorchScript file
torch_tensorrt.save(trt_gm, "trt.ts", output_format="torchscript", inputs=inputs)

import torch
import torch_tensorrt

inputs = [torch.randn((1, 3, 224, 224)).cuda()] # your inputs go here

# You can run this in a new python session!
model = torch.export.load("trt.ep").module()
# model = torch_tensorrt.load("trt.ep").module() # this also works
model(*inputs)

处理尺寸输入

默认情况下，可以在PyTorch模型中使用不同输入形状，并在运行时确定输出形状。然而，Torch-TensorRT作为一个提前编译（AOT）编译器，需要一些关于输入形状的先验信息来编译和优化模型。

使用torch.export处理动态形状（AOT）

在动态输入形状的情况下，必须提供(min_shape, opt_shape, max_shape)参数，以便模型可以针对这个输入形状范围进行优化。以下是静态和动态形状的示例用法。

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
# 使用静态形状编译
inputs = torch_tensorrt.Input(shape=[1, 3, 224, 224], dtype=torch.float32)
# 或使用动态形状编译
inputs = torch_tensorrt.Input(min_shape=[1, 3, 224, 224],
                              opt_shape=[4, 3, 224, 224],
                              max_shape=[8, 3, 224, 224],
                              dtype=torch.float32)
trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs)

自定义动态形状约束

给定输入x = torch_tensorrt.Input(min_shape, opt_shape, max_shape, dtype)，Torch-TensorRT尝试在torch.export跟踪期间自动设置约束，通过相应地构建torch.export.Dim对象来实现。有时，可能需要设置额外的约束，如果未指定这些约束，Torchdynamo会报错。如果需要为模型设置任何自定义约束（使用torch.export.Dim），建议先导出程序，然后再用Torch-TensorRT编译。

import torch
import torch_tensorrt

class MatMul(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, query, key):
        attn_weight = torch.matmul(query, key.transpose(-1, -2))
        return attn_weight

model = MatMul().eval().cuda()
inputs = [torch.randn(1, 12, 7, 64).cuda(), torch.randn(1, 12, 7, 64).cuda()]
seq_len = torch.export.Dim("seq_len", min=1, max=10)
dynamic_shapes=({2: seq_len}, {2: seq_len})
# 首先导出模型，自定义动态形状约束
exp_program = torch.export.export(model, tuple(inputs), dynamic_shapes=dynamic_shapes)
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs)
# 运行推理
trt_gm(*inputs)

使用torch.compile处理动态形状（JIT）

torch_tensorrt.compile(model, inputs, ir="torch_compile")返回一个配置为TensorRT后端的torch.compile封装函数。在ir=torch_compile的情况下，用户可以使用torch._dynamo.mark_dynamic API为输入提供动态形状信息，以避免TensorRT引擎的重新编译。

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = torch.randn((1, 3, 224, 224), dtype=float32)
# 这表明第0维是动态的，范围是[1, 8]
torch._dynamo.mark_dynamic(inputs, 0, min=1, max=8)
trt_gm = torch.compile(model, backend="tensorrt")
# 调用模型时发生编译
trt_gm(inputs)

# 修改批量大小时，不重新编译TRT引擎
inputs_bs2 = torch.randn((2, 3, 224, 224), dtype=torch.float32)
trt_gm(inputs_bs2)

进阶案例

案例1：ResNet

import torch
import torch_tensorrt
import torchvision.models as models


# 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()]

# 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,
)

# 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)

案例2：BERT

import torch
import torch_tensorrt
from transformers import BertModel

# Initialize model with float precision and sample inputs
model = BertModel.from_pretrained("bert-base-uncased").eval().to("cuda")
inputs = [
    torch.randint(0, 2, (1, 14), dtype=torch.int32).to("cuda"),
    torch.randint(0, 2, (1, 14), dtype=torch.int32).to("cuda"),
]

# Define backend compilation keyword arguments
compilation_kwargs = {
    "enabled_precisions": enabled_precisions,
    "debug": debug,
    "workspace_size": workspace_size,
    "min_block_size": min_block_size,
    "torch_executed_ops": torch_executed_ops,
}

# Build and compile the model with torch.compile, using Torch-TensorRT backend
optimized_model = torch.compile(
    model,
    backend="torch_tensorrt",
    dynamic=False,
    options=compilation_kwargs,
)
optimized_model(*inputs)

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。