最近闲来无事,还有三天回家过年,那么就利用这三天,来把tensorRT系统学习一下吧,我参考的是官方文档:https://docs.nvidia.com/deeplearning/tensorrt/index.html
1. 介绍:
本文档是希望能够帮助大家快速构建一个应用并且基于TensorRT engine来进行推理。TensorRT是一个可以保障高性能推理的SDK,TensorRT包含一个深度学习模型推理优化器,主要是为了给训练好的模型使用,同时还包含一个用于执行的runtime。
下面是使用TensorRT进行典型的深度学习开发的流程图:
2. 安装
安装有多种方法,有基于容器的(docker)、基于deb安装包的、基于pip的,具体可以看文档
官方推荐了一堆,而且文档写的云里雾里,我使用的是基于tar压缩文件的方式(Tar file Installation),其实就是在系统变量里面指定tensorRT的路径,教程使用的方式是在终端输入:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<TensorRT-${version}/lib>
但是为了方便我们可以直接在~/.bashrc
文件中进行修改,添加类似于下方的语句:
export LD_LIBRARY_PATH=/home/ubuntu/TensorRT-8.4.1.5/lib:$LD_LIBRARY_PATH
3. TensorRT 生态系统
TensorRT是一个又大又灵活的项目,可以处理各种转换和部署工作流,具体怎么用取决于你的任务。TensorRT对部署提供几个优化选项,但是所有的工作流都围绕着将你的模型转换为一个优化后的表征这样一个任务,这个表征在TensorRT中称之为engine。为你的模型构建一个TensorRT工作流涉及到选择正确的部署优化选项和为engin创建选择正确的参数组合。
你必须按照以下5个基本步骤来转换并部署你的模型:
TensorRT生态系统可以划分为以下两个类:
- 不同的用户使用不同路径来将他们的模型转换为优化的TensorRT engine。
- 不同的用户使用不同的runtime来部署他们的TensorRT engine。
下面是两种不同的方式示意图:
3.1 模型转换
主要有三种模型转换方式:
- 使用 TF-TRT
- 使用onnx文件自动转换
- 使用TensorRT API(C++ 或者 Python)来进行转换
对于TensorFlow模型的转换,TF-TRT同时提供模型转换和高级别的runtime API,并且对于TensorRT不支持的特定算子能够fall back到TensorFlow实现中去(感觉很牛的样子,但是我没用过)。
另一种方法是使用ONNX文件来进行转换,TensorRT支持从ONNX文件进行自动转换,要不是使用TensorRT API要不是使用trtexec工具,ONNX转换方式需要你模型中的所有操作必须被TensorRT支持(或者你必须对不支持的算子提供自定义的plug-ins实现),使用ONNX转换的结果是一个奇异的TensorRT engine(啥意思?没懂),比TF-TRT的开销更小。
对于极致的追求和更多自定义的可能性,你可以手动通过TensorRT API来进行模型构建,这种方法需要你一步一步的构建你的目标网络。在一个TensorRT网络被创建以后,你只要从模型导出模型权重并且加载到TensorRT 网络中去就好,可以参考:(这部分后面要好好看看,因为的确有不少网络的部署是自己写的engine workflow)
Creating A Network Definition From Scratch Using The C++ API
Creating A Network Definition From Scratch Using The Python API
3.2 模型部署
有3个使用TensorRT来部署模型的选项:
- 通过TensorFlow来进行
- 使用标准的TensorRT runtime API
- 使用NVIDIA Triton Inference Server
不同的方式决定了你转换模型的步骤,当使用TF-TRT的时候,对于部署最常见的选项是简单的部署在TensorFlow中,TF-TRT转换的结果是一个插入了TensorRT操作的TensorFlow图,这意味着你可以像你在使用TensorFlow其他模型一样运行TF-TRT模型。
TensorRT runtime API允许最低的开销和最细粒度的控制,但是对于TensorRT原生不支持的操作必须要通过plug-ins来实现(一些提前写好的plug-ins可以从这里获得,可以当做你自己写plug-in的参考资料)
还有一个NVIDIA Triton Inference Server,这个工具非常适合HTTP端去推理,详细可以去这里查看
3.3 选择正确的工作流
在选择如何转换和部署模型时候的两个最重要的因素是:
- 你选择的框架
- 你针对目标倾向的TensorRT runtime
下图是TensorRT设计到的流程图:
4. 使用ONNX的一个部署样例
我们转换一个预训练的ResNet-50模型,模型是ONNX格式的,ONNX格式的模型与框架无关,不管是TensorFlow还是Pytorch都可以转换成这种格式来进行部署。
首先通过wget来下载ResNet-50模型:
4.1 模型导出
可以参考 Exporting to ONNX from TensorFlow和 Exporting to ONNX from PyTorch.
这里直接使用现有的onnx啦:
wget https://download.onnxruntime.ai/onnx/models/resnet50.tar.gz
tar xzf resnet50.tar.gz
4.2 选择Batch Size
Batch Size对TensorRT的优化是非常重要的,一般来说,在推断中,当我们想要优先考虑延迟时,我们选择一个小的批处理大小,当我们想要优先考虑吞吐量时,我们选择一个更大的批处理大小。较大Batch size的处理时间较长,但减少了每个样本的平均处理时间。
如果你直到运行时才知道需要的批处理大小,这种情况下TensorRT能够动态地处理批大小(如何设置动态batch size,记得有一家面试问过我,我说我不知道)。也就是说,固定的batch size会让TensorRT进行额外的优化。动态的输入尺寸,参考dynamic shapes(这部分我们后面详细研究下,到底看看是啥东西),我们在导出ONNX的时候,已经设置了BATCH_SIZE=64
,所以这里不用单独设置,但是我们要记住这个数字,因为在部署的时候这个参数也是比较重要的。其他关于batch size的可以参考Batching(我们后面都会深入进去,先立个flag)
4.3 选择一个模型精度
这个对于涉及到部署的同学来说,极其重要。推理通常比训练要求更少的数值精度,在某些情况下,低精度能够在没有太多精度损失的前提下,给你更快的计算速度和更低的内存消耗。TensorRT支持TF32、FP32、FP16、INT8几个精度,精度问题可以查阅 Reduced Precision
FP32是很多框架的默认精度,我们也使用FP32:
import numpy as np
PRECISION = np.float32
4.4 转换模型
使用ONNX方式转换,有几种从ONNX转换到engine的方法,最常用的就是trtexec,运行如下命令:
trtexec --onnx=resnet50/model.onnx --saveEngine=resnet_engine.trt
--onnx
表示输入onnx路径, --saveEngine
表示保存位置
4.5 部署模型
在获得engine后,我们要决定如何来通过TensorRT使用这个engine,TensorRT运行时有两种类型:具有c++和Python绑定的独立runtime,以及与TensorFlow的原生集成。这节我们使用一个简化版的包装器ONNXClassifierWrapper
,我们称之为标准runtime。我们生成一个batch的随机dummy数据,并使用ONNXClassifierWrapper
来对该batch进行推理,更多runtime信息,参考 Understanding TensorRT Runtimes
注意,这里的onnx_helper并不是一个库,而是教程里面自己写的一个包装器。github文件在:https://github.com/NVIDIA/TensorRT/blob/master/quickstart/IntroNotebooks/onnx_helper.py
from onnx_helper import ONNXClassifierWrapper
import numpy as np
PRECISION = np.float32
N_CLASSES = 1000 # Our ResNet-50 is trained on a 1000 class ImageNet task
trt_model = ONNXClassifierWrapper("resnet_engine.trt", [BATCH_SIZE, N_CLASSES], target_dtype = PRECISION)
BATCH_SIZE=32
dummy_input_batch = np.zeros((BATCH_SIZE, 224, 224, 3), dtype = PRECISION)
predictions = trt_model.predict(dummy_input_batch)
注意教程中提到:wrapper不会加载和初始化engine,只有运行第一个batch的时候才会做这些,batching的更多资料,参考 Batching,TensorRT API参考 NVIDIA TensorRT API Reference。
5. TF-TRT Framework Integration
TF-TRT是一个Python的高阶接口,可以直接和TensorFlow模型做交互,允许你转换TensorFlow模型到TensorRT优化的模型,并且使用Python 高阶API来运行它。
TF-TRT同时提供转换路径和Python runtime,允许你运行任意的TensorFlow模型,TF-TRT有个很牛的地方啊,就是你可以同时转换模型中TensorRT支持的和不支持的layer,不需要额外创建plug-ins,感觉很牛有没有,毕竟自己写plug-ins很烦人,TF-TRT主要通过分析模型并且将子图传递给TensorRT,这样TensorRT engine就可能单独的进行转换。
但是这里需要注意的是,这种方式需要依赖整套的TensorFlow工具链,而不是像之前那种可以部署在单纯的Python和C++ runtime中的,这里的runtime被单独叫做TF-TRT Python runtime。从下面图片也能看粗来。对于tensorflow我也很久很久没有使用了,可以参考这个jupyter notebook。
6. ONNX转换和部署
感觉上面第四节已经讲过了,这里又更详细的说了一下,ONNX可以通过tf2onnx工具转换TensorFlow模型,这里我不去研究了,还是以pytorch为主吧,感兴趣直接看官网教程:https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html#export-from-tf
大概流程类似于下图:
跟着代码一起来实操一下呗:
-
从torchvision中导入ResNet-50模型的预训练权重
import torchvision.models as models resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
-
使用dummy数据,也就是假数据作为输入,来初始化onnx的输入
import torch BATCH_SIZE = 64 dummy_input=torch.randn(BATCH_SIZE, 3, 224, 224)
-
保存模型,其实对于导出模型有很多的 opset_version,这里我们后面也深入研究一下(还是立个flag)
import torch.onnx torch.onnx.export(resnext50_32x4d, dummy_input, "resnet50_onnx_model.onnx", verbose=False)
上面几步做完,就获得了一个.onnx
文件,后面就利用这个.onnx
文件来进行engine的生成吧。
有两种转换onnx到engine文件的方法,上文也有提到过,分别是:
- 使用trtexec
- 使用TensorRT API
这里我们使用trtexec来进行转换,指令如下:
trtexec --onnx=resnet50_onnx_model.onnx --saveEngine=resnet_engine.trt
这样就可以将resnet50_onnx_model.onnx
文件转换为resnet_engine.trt
这个engine文件,我们就成功获得一个engine文件啦。
还有另一个简单的方式就是使用上面提到过的ONNXClassifierWrappe
,快回头看吧!
7. 使用TensorRT Runtime API
TensorRT提供的标准C++和Python的runtime,比起TensorFlow要更优也更灵活,C++ API开销更小,但是Python的API更好用,因为Python有对应的data loaders还有类似于NumPy 和 SciPy这样的库,而且对于快速验证、调试和测试都更简单。
下面的例子是一个语义分割任务,使用TensorRT的C++和Python API。模型是ResNet-101作为backbone,接收任意大小的图片进行像素级别的预测。
步骤如下(又是在重复之前的):
- 设置test容器(container),通过Pytorch模型转换onnx并且使用trtexec转换成engine
- 使用C++ runtime来进行推理
- 使用Python runtime来进行推理
7.1 构建TensorRT engine(docker)
clone仓库,可以使用docker的方式(docker是另一种container工具,后面有机会我们也补上?),这里的docker你首先要安装docker,然后再运行,如果你本地没有,它会从远程拉取一个镜像下来。
git clone https://github.com/NVIDIA/TensorRT.git
cd TensorRT/quickstart
docker run --rm -it --gpus all -p 8888:8888 -v `pwd`:/workspace -w /workspace/SemanticSegmentation nvcr.io/nvidia/pytorch:20.12-py3 bash
python export.py
使用trtexec构建一个engine,其实这里我们上面已经操作过了。trtexec使用TensorRT ONNX parser来将ONNX模型加载到TensorRT的网络图中去,使用TensorRT的Builder API来生成优化后的engine。执行以下命令:
trtexec --onnx=fcn-resnet101.onnx --fp16 --workspace=64 --minShapes=input:1x3x256x256 --optShapes=input:1x3x1026x1282 --maxShapes=input:1x3x1440x2560 --buildOnly --saveEngine=fcn-resnet101.engine
trtexec构建engine还有很多可以配置的选项,参考NVIDIA TensorRT Developer Guide。你可以通过任意值的输入来进行校验:
trtexec --shapes=input:1x3x1026x1282 --loadEngine=fcn-resnet101.engine
其中–shapes表示对于推理,使用动态shape的input,如果成功会输出类似下面语句:
&&&& PASSED TensorRT.trtexec # trtexec --shapes=input:1x3x1026x1282 --loadEngine=fcn-resnet101.engine
7.2 使用C++方式来进行推理
编译并运行docker中的C++语义分割工程:
make
./bin/segmentation_tutorial
常规的C++推理步骤如下(参考 Deserializing A Plan ):
-
从文件中对engine进行反序列化,文件内容被加载到buffer中并且被反序列化到memory中去:
std::vector<char> engineData(fsize); engineFile.read(engineData.data(), fsize); std::unique_ptr<nvinfer1::IRuntime> runtime{nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger())}; std::unique_ptr<nvinfer1::ICudaEngine> mEngine(runtime->deserializeCudaEngine(engineData.data(), fsize, nullptr));
-
TensorRT执行上下文封装了执行状态,例如用于在推理期间保存中间激活张量的持久设备内存。
由于分割模型是在使用了动态形状的情况下构建的,因此必须为推理执行指定输入的形状。可以查询网络输出形状以确定输出buffer的相应size。(TensorRT支持按名称查询网络I/O绑定的索引)auto input_idx = mEngine->getBindingIndex("input"); assert(mEngine->getBindingDataType(input_idx) == nvinfer1::DataType::kFLOAT); auto input_dims = nvinfer1::Dims4{1, 3 /* channels */, height, width}; context->setBindingDimensions(input_idx, input_dims); auto input_size = util::getMemorySize(input_dims, sizeof(float)); auto output_idx = mEngine->getBindingIndex("output"); assert(mEngine->getBindingDataType(output_idx) == nvinfer1::DataType::kINT32); auto output_dims = context->getBindingDimensions(output_idx); auto output_size = util::getMemorySize(output_dims, sizeof(int32_t));
-
对于推理的准备,需要CUDA 设备的inputs和outputs memory已经分配好,图片被处理并且被拷贝到的memory中,且一系列的engine bindings已经被生成了(说实话,这个binding是啥我还真没弄清楚)。对于语义分割,输入数据使用均值为[0.485, 0.456, 0.406]和标准差为[0.229, 0.224, 0.225]进行了预处理,对于torchvision的模型预处理,参考这里,这里所有的操作使用RGBImageReader进行封装。
void* input_mem{nullptr}; cudaMalloc(&input_mem, input_size); void* output_mem{nullptr}; cudaMalloc(&output_mem, output_size); const std::vector<float> mean{0.485f, 0.456f, 0.406f}; const std::vector<float> stddev{0.229f, 0.224f, 0.225f}; auto input_image{util::RGBImageReader(input_filename, input_dims, mean, stddev)}; input_image.read(); auto input_buffer = input_image.process(); cudaMemcpyAsync(input_mem, input_buffer.get(), input_size, cudaMemcpyHostToDevice, stream);
-
使用上下文的executeV2或enqueueV2方法启动推理执行。执行完成后,我们将结果复制回主机缓冲区并释放所有设备内存分配。
void* bindings[] = {input_mem, output_mem}; bool status = context->enqueueV2(bindings, stream, nullptr); auto output_buffer = std::unique_ptr<int>{new int[output_size]}; cudaMemcpyAsync(output_buffer.get(), output_mem, output_size, cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream); cudaFree(input_mem); cudaFree(output_mem);
-
为了可视化结果,我们生成了一张伪图像,并将结果写入到output.ppm文件中,这些后处理工作被封装到
ArgmaxImageWriter
类中。(对于argmax是语义分割arm端部署的任务中非常烦人的事)const int num_classes{21}; const std::vector<int> palette{ (0x1 << 25) - 1, (0x1 << 15) - 1, (0x1 << 21) - 1}; auto output_image{util::ArgmaxImageWriter(output_filename, output_dims, palette, num_classes)}; output_image.process(output_buffer.get()); output_image.write();
argmax代码,非常有用,建议仔细研究
void ArgmaxImageWriter::process(const int* buffer) { mPPM.magic = "P6"; mPPM.w = mDims.d[3]; mPPM.h = mDims.d[2]; mPPM.max = 255; mPPM.buffer.resize(volume()); std::vector<std::vector<int>> colors; for (auto i = 0, max = mPPM.max; i < mNumClasses; i++) { std::vector<int> c{mPalette}; std::transform(c.begin(), c.end(), c.begin(), [i, max](int p){return (p*i) % max;}); colors.push_back(c); } for (int j = 0, HW = mPPM.h * mPPM.w; j < HW; ++j) { auto clsid{static_cast<uint8_t>(buffer[j])}; mPPM.buffer.data()[j*3] = colors[clsid][0]; mPPM.buffer.data()[j*3+1] = colors[clsid][1]; mPPM.buffer.data()[j*3+2] = colors[clsid][2]; } }
输入和输出如下:
7.3 使用Python来进行推理
首先安装pycuda:
pip install pycuda
然后启动jupyter notebook使用提供的token来加载log,打开地址为http://<host-ip-address>:8888
,你就直接在浏览器输入http://localhost:8888
应该就可以咯:
jupyter notebook --port=8888 --no-browser --ip=0.0.0.0 --allow-root
然后使用 tutorial-runtime.ipynb来完成后续步骤,流程和C++方式基本一致,我们也把它摘出来看一下:
-
导入模块及设置图片的输入输出
import numpy as np import os import pycuda.driver as cuda import pycuda.autoinit import tensorrt as trt import matplotlib.pyplot as plt from PIL import Image TRT_LOGGER = trt.Logger() # Filenames of TensorRT plan file and input/output images. engine_file = "fcn-resnet101.engine" input_file = "input.ppm" output_file = "output.ppm"
-
输入输出处理模块,分别是归一化和后处理,这里的归一化操作和可视化操作上面有讲解
# For torchvision models, input images are loaded in to a range of [0, 1] and # normalized using mean = [0.485, 0.456, 0.406] and stddev = [0.229, 0.224, 0.225]. def preprocess(image): # Mean normalization mean = np.array([0.485, 0.456, 0.406]).astype('float32') stddev = np.array([0.229, 0.224, 0.225]).astype('float32') data = (np.asarray(image).astype('float32') / float(255.0) - mean) / stddev # Switch from HWC to to CHW order return np.moveaxis(data, 2, 0) def postprocess(data): num_classes = 21 # create a color palette, selecting a color for each class palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1]) colors = np.array([palette*i%255 for i in range(num_classes)]).astype("uint8") # plot the segmentation predictions for 21 classes in different colors img = Image.fromarray(data.astype('uint8'), mode='P') img.putpalette(colors) return img
-
加载TensorRT engine,从文件中加载engine文件并且进行反序列化操作
def load_engine(engine_file_path): assert os.path.exists(engine_file_path) print("Reading engine from file {}".format(engine_file_path)) with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read())
-
推理的pipeline,这个流程还是比较重要的,基本就是TensorRT的工作顺序
TensorRT的推理pipeline包含如下步骤:
a. 创建一个execution contex并且指定输入形状(基于推理的图片纬度)
b. 为输入和输出分配CUDA device memory
c. 分配CUDA page-locked host memory高效地将输出拷贝出来(就是分配一种页锁定的主机内存,这样可以实现快速拷贝)
d. 使用host-to-device异步拷贝操作将处理好的图片信息转移到input memory中去
e. 使用异步的API来启动TensorRT推理
f. 使用device-to-host拷贝,将分割的输出结果转移到pagelocked host memory中去
g. 针对数据跨设备传输,一定要进行stream synchronization,也就是流的同步,这样可以保证推理操作可以正常结束,并且获得完整的数据
h. 可视化输出结果def infer(engine, input_file, output_file): print("Reading input image from file {}".format(input_file)) with Image.open(input_file) as img: input_image = preprocess(img) image_width = img.width image_height = img.height with engine.create_execution_context() as context: # Set input shape based on image dimensions for inference context.set_binding_shape(engine.get_binding_index("input"), (1, 3, image_height, image_width)) # Allocate host and device buffers bindings = [] for binding in engine: binding_idx = engine.get_binding_index(binding) size = trt.volume(context.get_binding_shape(binding_idx)) dtype = trt.nptype(engine.get_binding_dtype(binding)) if engine.binding_is_input(binding): input_buffer = np.ascontiguousarray(input_image) input_memory = cuda.mem_alloc(input_image.nbytes) bindings.append(int(input_memory)) else: output_buffer = cuda.pagelocked_empty(size, dtype) output_memory = cuda.mem_alloc(output_buffer.nbytes) bindings.append(int(output_memory)) stream = cuda.Stream() # Transfer input data to the GPU. cuda.memcpy_htod_async(input_memory, input_buffer, stream) # Run inference context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) # Transfer prediction output from the GPU. cuda.memcpy_dtoh_async(output_buffer, output_memory, stream) # Synchronize the stream stream.synchronize() with postprocess(np.reshape(output_buffer, (image_height, image_width))) as img: print("Writing output image to file {}".format(output_file)) img.convert('RGB').save(output_file, "PPM")
-
运行推理,这里面的函数在上面都有对应的定义
print("Running TensorRT inference for FCN-ResNet101") with load_engine(engine_file) as engine: infer(engine, input_file, output_file)
-
显示图片文章来源:https://www.toymoban.com/news/detail-828183.html
# 输入图片 plt.imshow(Image.open(input_file)) # 输出图片 plt.imshow(Image.open(output_file))
文章来源地址https://www.toymoban.com/news/detail-828183.html
到了这里,关于一、深入学习TensorRT,Getting Started篇的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!