如何评价NVIDIA发布的TRTorch?
关注者
293被浏览
87,272登录后你可以
不限量看优质回答私信答主深度交流精彩内容一键收藏
TRTorch 是 PyTorch -> TensorRT 模型转换器/编译器。原理和流程十分类似其他两款工具:TVM 和 Core 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: