一、TensorRT版本部署简介
模型训练好之后,就到了AI模型实际应用的最后一环 --- 模型部署。模型部署是将已经训练好的模型,部署到实际的应用场景中,所以模型部署更多是工程化的过程,就是解决实际的模型应用问题,本文采用NVIDIA推出的针对GPU卡的TensorRT对onnx模型转换部署。
二、模型压缩
模型压缩是对已经训练好的深度模型进行精简,进而得到一个轻量且准确率相当的网络,压缩后的网络具有更小的结构和更少的参数,可以有效降低计算和存储开销,便于部署在受限的硬件环境中。
训练的时候因为要保证前后向传播,每次梯度的更新是很微小的,这个时候需要相对较高的精度,一般来说需要float型,如FP32,32位的浮点型来处理数据,但是在推理(Inference)的时候,对精度的要求没有那么高,很多研究表明可以用低精度,如半长(16)的float,即FP16,也可以用8位的整型(INT8)来做推理(Inference)。所以,一般来说,在模型部署时会对模型进行压缩。模型压缩方法有:蒸馏,剪枝,量化等。
三、模型格式转换
首先研究员通过各种训练框架训练好的模型一般都需要进行模型格式适配,模型训练大家可以选择各种不同的训练框架,例如TensorFlow,Pytorch,PaddlePaddle,Caffe等等一系列的开源框架,这么多不同的训练框架他们训练出来的模型格式都有各自的标准,各不相同,部署要解决的第一个问题就是要适配各种不同的模型格式。但是如果要一个个训练框架去适配格式,工作量太大,也不适合扩展,所以微软联合Facebook等大厂推出一种中间格式ONNX,希望能解决多种模型格式适配的问题,就是无论是什么训练框架训练出来的模型格式,最终都是用ONNX格式来进行部署。所以一般模型部署可以跑的第一步要解决的问题就是模型格式转换。
1.初步了解部署内容
在软件工程中,部署指把开发完毕的软件投入使用的过程,包括环境配置、软件安装等步骤。类似地,对于深度学习模型来说,模型部署指让训练好的模型在特定环境中运行的过程。相比于软件部署,模型部署会面临更多的难题:
(1)运行模型所需的环境难以配置。深度学习模型通常是由一些框架编写,比如 PyTorch、TensorFlow。由于框架规模、依赖环境的限制,这些框架不适合在手机、开发板等生产环境中安装。
(2)深度学习模型的结构通常比较庞大,需要大量的算力才能满足实时运行的需求。模型的运行效率需要优化。
2.创建关于pytorch内容
用 PyTorch 实现一个超分辨率模型,并把模型部署到 ONNX Runtime 这个推理引擎上。首先,需要创建一个有 PyTorch 库的 Python 编程环境。如果你的 PyTorch 环境还没有装好,可以参考官方的入门教程。在这里推荐使用 conda 来管理 Python 库。
在介绍 ONNX 之前,先从本质上来认识一下神经网络的结构。神经网络实际上只是描述了数据计算的过程,其结构可以用计算图表示。比如 a+b 可,为了加速计算,一些框架会使用对神经网络“先编译,后执行”的静态图来描述网络。静态图的缺点是难以描述控制流(比如 if-else 分支语句和 for 循环语句),直接对其引入控制语句会导致产生不同的计算图。比如循环执行 n 次 a=a+b,对于不同的 n,会生成不同的计算图。
ONNX (Open Neural Network Exchange)是 Facebook 和微软在2017年共同发布的,用于标准描述计算图的一种格式。目前,在数家机构的共同维护下,ONNX 已经对接了多种深度学习框架和多种推理引擎。因此,ONNX 被当成了深度学习框架到推理引擎的桥梁,就像编译器的中间语言一样。由于各框架兼容性不一,通常只用 ONNX 表示更容易部署的静态图。
用下面的代码来把 PyTorch 的模型转换成 ONNX 格式的模型:
x = torch.randn(1, 3, 256, 256)
with torch.no_grad():
torch.onnx.export(
model,
x,
"srcnn.onnx",
opset_version=11,
input_names=['input'],
output_names=['output'])
其中,torch.onnx.export 是 PyTorch 自带的把模型转换成 ONNX 格式的函数。先看一下前三个必选参数:前三个参数分别是要转换的模型、模型的任意一组输入、导出的 ONNX 文件的文件名。从 PyTorch 的模型到 ONNX 的模型,本质上是一种语言上的翻译。直觉上的想法是像编译器一样彻底解析原模型的代码,记录所有控制流。但前面也讲到,通常只用 ONNX 记录不考虑控制流的静态图。因此,PyTorch 提供了一种叫做追踪(trace)的模型转换方法:给定一组输入,再实际执行一遍模型,即把这组输入对应的计算图记录下来,保存为 ONNX 格式。export 函数用的就是追踪导出方法,需要给任意一组输入,让模型跑起来。测试图片是三通道,256x256大小的,这里也构造一个同样形状的随机张量。
剩下的参数中,opset_version 表示 ONNX 算子集的版本。深度学习的发展会不断诞生新算子,为了支持这些新增的算子,ONNX会经常发布新的算子集,目前已经更新15个版本。假如令 opset_version = 11,即使用第11个 ONNX 算子集,是因为 SRCNN 中的 bicubic (双三次插值)在 opset11 中才得到支持。剩下的两个参数 input_names, output_names 是输入、输出 tensor 的名称。
如果上述代码运行成功,目录下会新增一个"model.onnx"的 ONNX 模型文件。就可以用下面的脚本来验证一下模型文件是否正确。
import onnx
onnx_model = onnx.load("model.onnx") try:
onnx.checker.check_model(onnx_model) except Exception:
print("Model incorrect") else:
print("Model correct")
其中,onnx.load 函数用于读取一个 ONNX 模型。onnx.checker.check_model 用于检查模型格式是否正确,如果有错误的话该函数会直接报错。该模型是正确的,控制台中应该会打印出"Model correct"。
接下来,让我们来看一看 ONNX 模型具体的结构是怎么样的。然后可以使用 Netron (开源的模型可视化工具)来可视化 ONNX 模型。把 srcnn.onnx 文件从本地的文件系统拖入网站,即可看到如下的可视化结果。
点击 input 或者 output,可以查看 ONNX 模型的基本信息,包括模型的版本信息,以及模型输入、输出的名称和数据类型。
(1)模型部署,指把训练好的模型在特定环境中运行的过程。模型部署要解决模型框架兼容性差和模型运行速度慢这两大问题。
(2)模型部署的常见流水线是“深度学习框架-中间表示-推理引擎”。其中比较常用的一个中间表示是 ONNX。
(3)深度学习模型实际上就是一个计算图。模型部署时通常把模型转换成静态的计算图,即没有控制流(分支语句、循环语句)的计算图。
PyTorch 框架自带对 ONNX 的支持,只需要构造一组随机的输入,并对模型调用。
3、模型部署中常见的难题
模型的动态化。出于性能的考虑,各推理框架都默认模型的输入形状、输出形状、结构是静态的。而为了让模型的泛用性更强,部署时需要在尽可能不影响原有逻辑的前提下,让模型的输入输出或是结构动态化。
新算子的实现。深度学习技术日新月异,提出新算子的速度往往快于 ONNX 维护者支持的速度。为了部署最新的模型,部署工程师往往需要自己在 ONNX 和推理引擎中支持新算子。
中间表示与推理引擎的兼容问题。由于各推理引擎的实现不同,对 ONNX 难以形成统一的支持。为了确保模型在不同的推理引擎中有同样的运行效果,部署工程师往往得为某个推理引擎定制模型代码,这为模型部署引入了许多工作量。
实现动态放大的超分辨率模型:
在原来的 SRCNN 中,图片的放大比例是写死在模型里的:
class SuperResolutionNet(nn.Module):
def __init__(self, upscale_factor):
super().__init__()
self.upscale_factor = upscale_factor
self.img_upsampler = nn.Upsample(
scale_factor=self.upscale_factor,
mode='bicubic',
align_corners=False)
...
def init_torch_model():
torch_model = SuperResolutionNet(upscale_factor=3)
使用 upscale_factor 来控制模型的放大比例。初始化模型的时候,假如默认令 upscale_factor 为 3,生成了一个放大 3 倍的 PyTorch 模型。这个 PyTorch 模型最终被转换成了 ONNX 格式的模型。这样需要一个放大 4 倍的模型,需要重新生成一遍模型,再做一次到 ONNX 的转换。
现在,假设要做一个超分辨率的应用。此时用户希望图片的放大倍数能够自由设置。在交给用户的时候,只有一个 .onnx 文件和运行超分辨率模型的应用程序。我们在不修改 .onnx 文件的前提下改变放大倍数。
因此,这时必须修改原来的模型,令模型的放大倍数变成推理时的输入。
import torch
from torch import nn
from torch.nn.functional import interpolate
import torch.onnx
import cv2
import numpy as np
class SuperResolutionNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
self.relu = nn.ReLU()
def forward(self, x, upscale_factor):
x = interpolate(x,
scale_factor=upscale_factor,
mode='bicubic',
align_corners=False)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.conv3(out)
return out
def init_torch_model():
torch_model = SuperResolutionNet()
state_dict = torch.load('srcnn.pth')['state_dict']
# Adapt the checkpoint
for old_key in list(state_dict.keys()):
new_key = '.'.join(old_key.split('.')[1:])
state_dict[new_key] = state_dict.pop(old_key)
torch_model.load_state_dict(state_dict)
torch_model.eval()
return torch_model
model = init_torch_model()
input_img = cv2.imread('face.png').astype(np.float32)
# HWC to NCHW input_img = np.transpose(input_img, [2, 0, 1]) input_img = np.expand_dims(input_img, 0)
# Inference torch_output = model(torch.from_numpy(input_img), 3).detach().numpy()
# NCHW to HWC torch_output = np.squeeze(torch_output, 0) torch_output = np.clip(torch_output, 0, 255) torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8)
# Show image cv2.imwrite("face_torch_2.png", torch_output)
SuperResolutionNet 未修改之前,nn.Upsample 在初始化阶段固化了放大倍数,而 PyTorch 的 interpolate 插值算子可以在运行阶段选择放大倍数。因此,在新脚本中使用 interpolate 代替 nn.Upsample,从而让模型支持动态放大倍数的超分。 在使用模型推理时,我们把放大倍数设置为 3。最后,图片保存在文件 "face_torch_2.png" 中。一切正常的话,"face_torch_2.png" 和 "face_torch.png" 的内容一模一样。
通过简单的修改,PyTorch 模型已经支持了动态分辨率。现在来尝试一下导出模型:
x = torch.randn(1, 3, 256, 256)
with torch.no_grad():
torch.onnx.export(model, (x, 3),
"srcnn2.onnx",
opset_version=11,
input_names=['input', 'factor'],
output_names=['output'])
- tensorrt模型转换和部署
1、tensorrt模型转换
(1)当batch=1时并用tensorrt:22.02-py3进行模型转换:
sudo docker run --gpus all -v $(pwd):/work -it nvcr.io/nvidia/tensorrt:22.02-py3 bash
cd /work
trtexec --onnx=swin_base_new3.onnx --saveEngine=model.plan --workspace=1024 --explicitBatch
(2)当batch=-1时并用tensorrt:22.02-py3进行模型转换:
sudo docker run --gpus all -v $(pwd):/work -it registry.cn-guangzhou.aliyuncs.com/nvidia-images/tensorrt:21.06-py3-opencv bash
cd /work
trtexec --onnx=scene_0718.onnx --saveEngine=scene.plan --workspace=1024 --minShapes=input:1x3x224x224 --optShapes=input:8x3x224x224 --maxShapes=input:16x3x224x224 --explicitBatch
(3)模型是yolov5时需要对应yolov5相关cpp文件和对应脚本转换:
#!/bin/bash
CURRENT_PATH=$(pwd)
MODEL_FILE=helmet_detect_0701.pt
MODEL_NAME=yolov5m #s,m,l,x
WTS_FILE=${MODEL_NAME}.wts
ENGINE_FILE=${MODEL_NAME}.plan
#sudo docker run --gpus all \
# -v ${CURRENT_PATH}:/work registry.cn-guangzhou.aliyuncs.com/nvidia-images/yolov5:4.0 \
# python3 gen_wts.py --model=/work/${MODEL_FILE} --wts=/work/${WTS_FILE}
sudo docker run --gpus all \
-v ${CURRENT_PATH}:/work \
-w /work \
-it registry.cn-guangzhou.aliyuncs.com/nvidia-images/tensorrt:21.06-py3-opencv \ bash -c 'cd yolov5-4.0-nms-helmet && bash run.sh'
其中21.06-py3-opencv或者22.02-py3-opencv需要在tensorrt:21.06-py3和tensorrt:22.02-py3容器里增加c++版本opencv即可。
2、模型部署
HTTP_PORT=7020
MODEL_DIR=$(dirname $(pwd))/models
DOCKER_IMAGE=nvcr.io/nvidia/tritonserver:22.02-py3
LOG_VERBOSE=1
NAME="ocr-docker"
sudo docker run \
--name ${NAME} \
--gpus '"device=1"' \
--rm --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 \
--network=host \
-v${MODEL_DIR}:/models \
${DOCKER_IMAGE} \
tritonserver --model-repository=/models \
--http-port=${HTTP_PORT} \
--allow-grpc=false \
--allow-metrics=false \
--allow-sagemaker=false \
--log-verbose=1
Docker启动服务脚本,值得注意的是模型tensorrt版本与tritonserver服务版本需要一致。
3、模型推理和前后处理
(1)前处理:因为模型推理的输入是Tensor(多维矩阵)数据,但是正常AI应用的输入都是图片,视频,文字等数据,所以前处理就是要将业务的输入数据(图像,视频,文字等)预先处理成模型推理可以接收的数据---Tensor(多维矩阵)。以图像处理为例,前处理动作就包括但不限于:图像格式转换,颜色空间变换,图像变换(resize,warpaffine,图像滤波等操作。
OpenCV就是intel推出开源的跨平台的计算机视觉库。
在预处理过程中为了提高处理效率,预处理可以选择在GPU上处理,一般选用DALI进行处理,下面是某个模型预处理DALI脚本:
import nvidia.dali.fn as fn
import nvidia.dali as dali
import nvidia.dali.types as types
import cv2
import numpy as np
import time
def get_imgs():
image_data = np.fromfile("5.jpg", dtype="uint8")
data = [image_data for i in range(8)]
return data
input_dim = [640,640]
@dali.pipeline_def(batch_size=32, num_threads=16, device_id=0)
def pipe():
data = fn.external_source(device="cpu", name="DALI_INPUT_0", source=get_imgs)
raw_images = fn.decoders.image(data, device="mixed", output_type=types.RGB)
images = fn.resize(raw_images, size=input_dim, mode="not_larger", max_size=input_dim)
images = fn.pad(images, fill_value=128, shape=[input_dim[0], input_dim[1], 3])
images = fn.transpose(images, perm=[2, 0, 1])
images = fn.cast(images, dtype=types.FLOAT)
return images
def main():
pipe().serialize(filename='model.dali')
print("serialized!")
pipeline = pipe()
pipeline.build()
o = pipeline.run()
image = np.transpose(np.array(o[0][0].as_cpu()),[1,2,0])[:,:,::-1]
cv2.imwrite("out0.jpg",image)
if __name__ == '__main__':
main()
其中DALI脚本对输入data进行解码然后对图像数据resize以及填充等处理并输出处理结果,此处DALI模型一般与推理模型级联部署处理。
(2)模型推理:模型推理应该是模型部署pipline中最核心的部分。就是需要在实际应用环境中(具体部署设备)将实际输入的数据(转换成Tensor数据后)在训练好的模型中跑通,并且性能和精度等商业指标上达到预期效果。这个过程包括了对部署设备的适配(CPU/GPU/DSP/NPU),要想将模型跑在任何一种设备上,都需要提前针对设备进行适配,并且还要保证性能和精度等指标。本文采用tensorrt对onnx模型转换和tritonserver服务部署(NVIDIA推出的针对GPU卡的TensorRT)。
市面上有非常多的开源深度学习推理框架都是在解决模型推理相关的问题。例如:国内各大厂推出的开源的推理框架:OpenPPL、NCNN、TNN、MNN、PaddleLite、Tengine等等,还有intel针对intel芯片的OpenVINO等。
后处理:就是将模型推理后的Tensor数据转换成业务可以识别的特征数据(不同的业务会呈现不同的最终效果数据)。