《torch2trt》
模型的 accuracy 和 speed 是鱼和熊掌不可兼得,所以如何进行有效的模型加速就很重要。之前尝试过模型剪枝、量化和蒸馏的方法,发现除了蒸馏可以提升小模型的效果外,其余的效果大都停留在学术层面上,实际用起来的模型效果影响还是比较明显的。直接结合硬件层面对模型的加速是比较合理的,所以模型直接转 TensorRT 加速一直都是业界用的比较多方案,本文介绍一个 NVIDIA 开源的 torch2trt,目前用起来还是有很多的 bug,但是可以自定义 converter 也不至于遇到问题束手无策。之前转 TRT 都是 caffe/tf/torch -> onnx -> trt 这个链路的,中间还需要编译 TensorRT。而直接用 NVIDIA 的这个代码就比较方便一些。
Key Words:torch2trt
Beijing, 2022
https://github.com/NVIDIA-AI-IOT/torch2trt
以官方的 resnet50 为例的转换程序如下:
import time
import torch
import numpy as np
import tensorrt as trt
from torch2trt import torch2trt
from torch2trt import tensorrt_converter
from torchvision.models.resnet import resnet50
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# create some regular pytorch model...
model = resnet50(pretrained=True).eval().cuda()
# convert to TensorRT feeding sample data as input
start = time.time()
model_trt = torch2trt(model, [x], fp16_mode=True, strict_type_constraints=True, use_onnx=True)
print("Convert cost time: %s"%((time.time() - start)*1000))
安装 torch2trt
基本上没什么坑,当时遇到了点小问题,我在 issue 里面也放了解决的评论
先对比 trt fp32/16 和 pytorch 的diff
np.testing.assert_almost_equal(model_trt(x).data.cpu().numpy(), model(x).data.cpu().numpy(), decimal=2)
fp16 的 diff 还是挺大的,fp32还好
对比 trt fp32/fp16 和 pytorch 的性能
def test_latency(model, model_trt, x):
### warm_up
model(x)
model(x)
model(x)
start = time.time()
for i in range(50):
model(x)
torch.cuda.synchronize() # 同步 gpu 操作结束
print("Torch Cost time: %s ms"%((time.time() - start) * 1000))
### warm_up
model_trt(x)
model_trt(x)
model_trt(x)
start = time.time()
for i in range(50):
model_trt(x)
torch.cuda.synchronize()
print("TRT Cost time: %s ms"%((time.time() - start) * 1000))
test_latency(model, model_trt, x)
不同型号的 GPU 不一样,就不放我的结果了
存储和加载 trt 模型
torch2trt模型保存并加载
torch.save(net_trt.state_dict(), 'resnet50_trt.pth')
model_trt = torch2trt.TRTModule()
model_trt.load_state_dict(torch.load('resnet50_trt.pth'))