vlambda博客
学习文章列表

什么是 Torch-TensorRT

Torch-TensorRT 的工作原理

• 降低 TorchScript 模块
• 转换
• 执行

降低 TorchScript 模块

转换

• 具有静态值的节点被评估并映射到常量。
• 描述张量计算的节点被转换为一个或多个 TensorRT 层。
• 其余节点保留在 TorchScripting 中，形成一个混合图，作为标准 TorchScript 模块返回。

执行

Torch-TensorRT 功能

支持 INT8

• 训练后量化 (PTQ)
• 量化感知训练 (QAT)

稀疏性

• TensorRT 支持在这些 Tensor Core 上注册和执行一些稀疏层的深度学习模型。
• Torch-TensorRT 扩展了对卷积和全连接层的支持。

示例：图像分类的吞吐量比较

安装和先决条件

• 具有 NVIDIA GPU、计算架构 7 或更早版本的 Linux 机器
• 已安装 Docker，19.03 或更高版本
• 一个 Docker 容器，包含 PyTorch、Torch-TensorRT 以及从NGC 目录中提取的所有依赖项

Jupyter lab --allow-root --IP=0.0.0.0 --NotebookApp.token=’TensorRT’ --port 8888

在浏览器上连接到 JupyterLab 的图形用户界面后，可以创建一个新的 Jupyter 笔记本。首先安装timm，这是一个包含预训练计算机视觉模型、权重和脚本的 PyTorch 库。EfficientNet-b0从这个库中拉出模型。

pip install timm

导入相关库并nn.Module为EfficientNet-b0.

import torch
import torch_tensorrt
import timm
import time
import numpy as np
import torch.backends.cudnn as cudnn

torch.hub._validate_not_a_forked_repo=lambda a,b,c: True

efficientnet_b0 = timm.create_model('efficientnet_b0',pretrained=True)

通过将随机浮点数的张量传递给该对象的forward方法，您可以从此模型中获得预测。efficientnet_b0

model = efficientnet_b0.eval().to("cuda")
detections_batch = model(torch.randn(128, 3, 224, 224).to("cuda"))
detections_batch.shape

这将返回一个 [128, 1000] 的张量，对应于 128 个样本和 1,000 个类。

要通过 PyTorch JIT 和 Torch-TensorRT AOT 编译方法对该模型进行基准测试，编写一个简单的基准实用程序函数：

cudnn.benchmark = True

def benchmark(model, input_shape=(1024, 3, 512, 512), dtype='fp32', nwarmup=50, nruns=1000):
    input_data = torch.randn(input_shape)
    input_data = input_data.to("cuda")
    if dtype=='fp16':
        input_data = input_data.half()
        
    print("Warm up ...")
    with torch.no_grad():
        for _ in range(nwarmup):
            features = model(input_data)
    torch.cuda.synchronize()
    print("Start timing ...")
    timings = []
    with torch.no_grad():
        for i in range(1, nruns+1):
            start_time = time.time()
            pred_loc  = model(input_data)
            torch.cuda.synchronize()
            end_time = time.time()
            timings.append(end_time - start_time)
            if i%10==0:
                print('Iteration %d/%d, avg batch time %.2f ms'%(i, nruns, np.mean(timings)*1000))

    print("Input shape:", input_data.size())
    print('Average throughput: %.2f images/second'%(input_shape[0]/np.mean(timings)))

现在已准备好对此模型执行推理。

使用 PyTorch 和 TorchScript 进行推理

首先，采用 PyTorch 模型并计算批量大小为 1 的平均吞吐量：

model = efficientnet_b0.eval().to("cuda")
benchmark(model, input_shape=(1, 3, 224, 224), nruns=100)

可以使用 TorchScript JIT 模块重复相同的步骤：

traced_model = torch.jit.trace(model, torch.randn((1,3,224,224)).to("cuda")])
torch.jit.save(traced_model, "efficientnet_b0_traced.jit.pt")
benchmark(traced_model, input_shape=(1, 3, 224, 224), nruns=100)

PyTorch 和 TorchScript JIT 报告的平均吞吐量将是相似的。

使用 Torch-TensorRT 进行推理

要使用 Torch-TensorRT 以混合精度编译模型，请运行以下命令：

trt_model = torch_tensorrt.compile(model, 
    inputs= [torch_tensorrt.Input((1, 3, 224, 224))],
    enabled_precisions= { torch_tensorrt.dtype.half} # Run with FP16
)

最后，对这个 Torch-TensorRT 优化模型进行基准测试：

benchmark(trt_model, input_shape=(1, 3, 224, 224), nruns=100, dtype="fp16")

基准测试结果

这是在批量大小为 1 的 NVIDIA A100 GPU 上取得的结果。

图 6. 在批量大小为 1 的 NVIDIA A100 GPU 上比较原生 PyTorch 与 Torch-TensorRt 的吞吐量

只需一行代码进行优化，Torch-TensorRT 即可将模型性能加速高达 6 倍。它确保了 NVIDIA GPU 的最高性能，同时保持了 PyTorch 的易用性和灵活性。

Ashish Sardana