如何评价NVIDIA发布的TRTorch?

https://github.com/NVIDIA/TRTorch 英伟达发布了基于Pytorch的编译器,号称可以让你无感的运行JIT的模型但是是在…
关注者
293
被浏览
87,272
登录后你可以
不限量看优质回答私信答主深度交流精彩内容一键收藏

TRTorchPyTorch -> TensorRT 模型转换器/编译器。原理和流程十分类似其他两款工具:TVMCore ML除了 TRTorch 只针对 NVIDIA GPU。

工作流程大致如下:

- 训练 / 下载 -> 源 PyTorch 模型
- PyTorch 模型 -> TorchScript 模型: torch.jit.script()/trace()
- TorchScript 模型 -> TensorRT 模型: trtorch.compile()

对应到官方的 Python 小栗子 :

import torch
import torchvision
import trtorch

# 训练 / 下载 -> 源 PyTorch 模型
model = torchvision.models.alexnet(pretrained=True).eval().cuda()

# PyTorch 模型 -> TorchScript 模型
data = torch.randn((1, 3, 224, 224)).to("cuda")
traced_model = torch.jit.trace(model, [data])

# TorchScript 模型 -> TensorRT 模型
compiled_trt_model = trtorch.compile(traced_model, {
    "input_shapes": [data.shape],
    "op_precision": torch.half, # Run in FP16
})

# 转换过的模型可以获得 NVIDIA GPU 加速
results = compiled_trt_model(data.half())

理论上如果采用 NVIDIA 推理部署方案,使用 TRTorch 转换后的模型可以获得针对 NVIDIA GPU 硬件的额外优化。根据官方 GTC 2020 Talk 给出的 ResNet50 性能对比(下图),相对直接使用 TensorRT 基本模型转换没有 overhead:

Reference: