Pytorch导出onnx模型,C++转化为TensorRT并实现推理过程

Pytorch导出onnx模型,C++转化为TensorRT并实现推理过程

前言

  1. 本文为旨在实现整个Python导出PyTorch模型,C++转化为TensorRT并实现推理过程过程,只与模型推理,模型部署相关,不涉及模型训练。
  2. 为突出整个部署过程而非具体模型本身,本文模型就采用最简单的分类模型,并且直接使用 torchvision.model 中的权重。检测、分割等其他模型在前后处理部分会有不同,但是模型本身的导出、转换和推理的过程基本是一致的。
  3. 本文会先用 Pytorch 测试一个分类模型在一张测试图片上的结果,将其转换为 onnx 模型,再用 onnxruntime 测试结果,再用 C++ 将其转换为 TensorRT 模型,再测试推理结果。预期三者测试结果一致,则转换成功。
  4. 如果想要测试速度,Python 可以使用 time.perf_counter(), C++ 可以使用 std::chrono::high_resolution_clock 。建议多测数据集中的一些图片,计算推理时间的均值和方差,而不是只测一张图片。

1 Pytorch模型推理测试导出onnx

这部分我们使用 torchvision 实例化一个简单的 ResNet50 分类模型,并将其导出为 onnx 模型。在这个过程中,我们还需要使用一张图片进行推理,并记录下 Python 模型的输出,方便我们后面到处 TensoRT 模型并进行推理时进行准确性的验证。

由于 torchvision 中的 resnet50 分类模型中是没有进行最后的 softmax 操作的,这里我们为了之后使用方便,自己新建一个类 ResNet50_wSoftmax 将后处理 softmax 添加到模型中一起导出。

这也是 pytorch 导出 onnx 模型的一个推荐的方式,就是将一些必要后处理添加到模型中一起导出,这样做有两个优点:

  • 可以直接得到端到端的 onnx/tensorrt 模型,不必在外面再做后处理操作
  • 再之后我们会将 onnx 模型转换为 tensorrt 模型,在转换过程中 tensorrt 会对我们的模型进行一些针对特定的 Nvidia GPU 的推理优化,我们将后处理一起合并到 onnx 模型中,可能可以使得一些算子操作再转换为 tensorrt 的过程中同样得到优化。

最终代码如下:

# export_onnx.py
import torch
import torchvision.models as models
import cv2
import numpy as np

class ResNet50_wSoftmax(torch.nn.Module):
    # 将softmax后处理合并到模型中,一起导出为onnx
    def __init__(self):
        super().__init__()
        self.base_model = models.resnet50(pretrained=True)
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        y = self.base_model(x)
        prob = self.softmax(y)
        return prob

def preprocessing(img):
    # 预处理:BGR->RGB、归一化/除均值减标准差
    IMAGENET_MEAN = [0.485, 0.456, 0.406]
    IMAGENET_STD = [0.229, 0.224, 0.225]
    img = img[:, :, ::-1]
    img = cv2.resize(img, (224, 224))
    img = img / 255.0
    img = (img - IMAGENET_MEAN) / IMAGENET_STD
    img = img.transpose(2, 0, 1).astype(np.float32)
    tensor_img = torch.from_numpy(img)[None]
    return tensor_img

if __name__ == '__main__':
    # model = models.resnet50(pretrained=True)
    image_path = 'test.jpg'
    img = cv2.imread(image_path)
    tensor_img = preprocessing(img)
    model = ResNet50_wSoftmax()   # 将后处理添加到模型中
    model.eval()
    pred = model(tensor_img)[0]
    max_idx = torch.argmax(pred)
    print(f"test_image: {image_path}, max_idx: {max_idx}, max_logit: {pred[max_idx].item()}")

    dummpy_input = torch.zeros(1, 3, 224, 224)  # onnx的导出需要指定一个输入,这里直接用上面的tenosr_img也可
    torch.onnx.export(
            model, dummpy_input, 'resnet50_wSoftmax.onnx',
            input_names=['image'],
            output_names=['predict'],
            opset_version=11,
            dynamic_axes={'image': {0: 'batch'}, 'predict': {0: 'batch'}}		# 注意这里指定batchsize是动态可变的
    )

执行结果会输出:

test_image: test.jpg, max_idx: 971, probability: 0.994541585445404

这些结果我们一会测试 onnx/tensorrt 模型时用于比对转换是否有误差。并得到一个 onnx 模型文件:classifier.onnx

2 onnxruntime推理测试

我们将刚刚得到的 classifier.onnx ,用 onnxruntime 来进行推理测试,看结果是否相同。

这里,我们就复用刚才测试 pytorch 模型时的预处理函数,整个 onnxruntime 推理测试代码如下:

import onnxruntime as ort
import numpy as np
import cv2
from export_onnx import preprocessing

image_path = 'test.jpg'
ort_session = ort.InferenceSession("classifier.onnx") # 创建一个推理session

img = cv2.imread(image_path)
input_img = preprocessing(img)[None]

pred = ort_session.run(None, { 'image' : input_img } )[0][0]
max_idx = np.argmax(pred)
print(f"test_image: {image_path}, max_idx: {max_idx}, probability: {pred[max_idx]}")

输出:

test_image: test.jpg, max_idx: 971, probability: 0.994541585445404

可以看到,跟我们 pytorch 模型的测试结果是一致的。

3 C++ onnx模型转换为tensorrt模型

本部分重度参考自课程:tensorRT从零起步迈向高性能工业级部署(就业导向)

我们进行模型部署推理肯定是追求极致的推理速度,这时再用 Python 来进行转换和推理就不合适了,接下来我们就转战到 C++ 上,将onnx模型转换为tensorrt模型。

对于大部分深度学习部署的 C/C++ 的初学者而言,环境配置都是个老大难的问题。本身 C/C++ 的包管理就不如 Python 的 pip、conda 等来的直接方便,再加上各种 nvidia driver/cuda/cudnn/cuda-runtime 的各种版本不对齐的问题,包括笔者在内的许多萌新们初期总是会在环境配置遇到许多问题。但是本文关注的重点是整个模型转换和部署的过程,不可能花大篇幅再去介绍环境配置,将来有机会再单独写一篇介绍 Python/C++ 深度学习模型部署时环境配置的问题,这里就直接给出笔者使用的关键软硬件的版本号/型号。

GPU: RTX 3060ti 12GB

OS: ubuntu 18.04

gcc: 7.5

TensorRT: 8.x

CUDA: 11.2

cuDNN: 8.x

头文件

包含的头文件:

// tensorrt相关
#include 
#include 

// onnx解析器相关
#include 

// cuda_runtime相关
#include 

// 常用头文件
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

// opencv
#include 

logger类

首先我们要准备一个 logger 类,来打印构建 tensorrt 模型过程中的一些错误或警告。按照指定的严重性程度 (severity),来打印信息。

inline const char* severity_string(nvinfer1::ILogger::Severity t) {
	switch (t) {
		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
		case nvinfer1::ILogger::Severity::kERROR: return "error";
		case nvinfer1::ILogger::Severity::kWARNING: return "warning";
		case nvinfer1::ILogger::Severity::kINFO: return "info";
		case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
		default: return "unknown";
	}
}

class TRTLogger : public nvinfer1::ILogger {
public:
	virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {
		if (severity <= Severity::kWARNING) {
			if (severity == Severity::kWARNING) printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
			else if (severity == Severity::kERROR) printf("\031[33m%s: %s\033[0m\n", severity_string(severity), msg);
			else printf("%s: %s\n", severity_string(severity), msg);
		}
	}
};

build_model函数

build_model 函数,各步骤已在代码中添加注释:

bool build_model() {
	if (isFileExist( "classifier.trtmodel" )) {
		printf("classifier.trtmodel already exists.\n");
		return true;
	}

	TRTLogger logger;
    
  // 下面的builder, config, network是基本需要的组件
  // 形象的理解是你需要一个builder去build这个网络,网络自身有结构,这个结构可以有不同的配置
	nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
  // 创建一个构建配置,指定TensorRT应该如何优化模型,tensorRT生成的模型只能在特定配置下运行
	nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
  // 创建网络定义,其中createNetworkV2(1)表示采用显性batch size,新版tensorRT(>=7.0)时,不建议采用0非显性batch size
	nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

  // onnx parser解析器来解析onnx模型
	auto parser = nvonnxparser::createParser(*network, logger);
	if (!parser->parseFromFile("classifier.onnx", 1)) {
		printf("Failed to parse classifier.onnx.\n");
		return false;
	}
	
  // 设置工作区大小
	printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
	config->setMaxWorkspaceSize(1 << 28);
  
  // 需要通过profile来使得batchsize时动态可变的,这与我们之前导出onnx指定的动态batchsize是对应的
	int maxBatchSize = 10;
	auto profile = builder->createOptimizationProfile();
	auto input_tensor = network->getInput(0);
	auto input_dims = input_tensor->getDimensions();

  // 设置batchsize的最大/最小/最优值
	input_dims.d[0] = 1;
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);

	input_dims.d[0] = maxBatchSize;
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);
	config->addOptimizationProfile(profile);

  // 开始构建tensorrt模型engine
	nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  
	if (engine == nullptr) {
		printf("Build engine failed.\n");
		return false;
	}
	
  // 将构建好的tensorrt模型engine反序列化(保存成文件)
	nvinfer1::IHostMemory* model_data = engine->serialize();
	FILE* f = fopen("classifier.trtmodel", "wb");
	fwrite(model_data->data(), 1, model_data->size(), f);
	fclose(f);

  // 逆序destory掉指针
	model_data->destroy();
	engine->destroy();
	network->destroy();
	config->destroy();
	builder->destroy();

	printf("Build Done.\n");
	return true;
}

调用 build_model 函数成功后,我们会得到一个 classifier.trtmodel 文件。

make_nvshared

上面的实现有个比较不优雅的地方,对于我们创建的 builderconfig 等指针,我们都需要一一进行 destroy,从而避免内存泄漏。实际上,这里我们可以通过共享指针,来实现自动释放。

shared_ptr<_T> make_nvshared(_T *ptr) {
	return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}

在这里指定一下释放内存的方式,之后就可以通过类似:

auto network = make_nvshared(builder->createNetworkV2(1));

这样的方式创建智能指针,他会自己 destroy 释放,这样最后几行 destory 就不用写了。

4 tensorrt模型推理测试

我们上一步已经成功将 onnx 模型导出为了 tensorrt 模型,现在我们用 tensorrt 模型来进行推理,看一下结果是否与之前 pytorch 和 onnx 推理的结果一致,如果一致,则模型转换成功。

load_file

load_file 函数用于加载我们的 tensorrt 模型:

vector<unsigned char> load_file(const string& file) {
	ifstream in(file, ios::in | ios::binary);
	if (!in.is_open()) return {};

	in.seekg(0, ios::end);
	size_t length = in.tellg();

	vector<uint8_t> data;
	if (length > 0) {
		in.seekg(0, ios::beg);
		data.resize(length);

		in.read((char*)&data[0], length);
	}
	in.close();
	return data;
}

inference

void inference(const string& image_path) {
	TRTLogger logger;
  // 加载模型
	auto engine_data = load_file("classifier.trtmodel");
  // 执行推理前,需要创建一个推理的runtime接口实例。与builer一样,runtime需要logger
	auto runtime = make_nvshared(nvinfer1::createInferRuntime(logger));
	auto engine = make_nvshared(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));
	if (engine == nullptr) {
		printf("Deserialize cuda engine failed.\n");
		runtime->destroy();
		return;
	}

	if (engine->getNbBindings() != 2) {
		printf("Must be single input, single Output, got %d output.\n", engine->getNbBindings() - 1);
		return;
	}

  // 创建CUDA流,以确定这个batch的推理是独立的
	cudaStream_t stream = nullptr;
	checkRuntime(cudaStreamCreate(&stream));
	auto execution_context = make_nvshared(engine->createExecutionContext());

	int input_batch = 1;
	int input_channel = 3;
	int input_height = 224;
	int input_width = 224;

  // 准备好input_data_host和input_data_device,分别表示内存中的数据指针和显存中的数据指针
  // 一会儿将预处理过的图像数据搬运到GPU
	int input_numel = input_batch * input_channel * input_height * input_width;
	float* input_data_host = nullptr;
	float* input_data_device = nullptr;

	checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));

  // 图片读取与预处理,与之前python中的预处理方式一致:
  // BGR->RGB、归一化/除均值减标准差
  float mean[] = {0.406, 0.456, 0.485};
	float std[] = {0.225, 0.224, 0.229};
  
	auto image = cv::imread(image_path);
	cv::resize(image, image, cv::Size(input_width, input_height));

	int image_area = image.cols * image.rows;
	unsigned char* pimage = image.data;
	float* phost_b = input_data_host + image_area * 0;
	float* phost_g = input_data_host + image_area * 1;
	float* phost_r = input_data_host + image_area * 2;
	for (int i=0; i<image_area; ++i, pimage += 3) {
		 *phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];
		 *phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];
		 *phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];
	 }

  // 进行推理
	checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel *sizeof(float), cudaMemcpyHostToDevice, stream));

	const int num_classes = 1000;
	float output_data_host[num_classes];
	float* output_data_device = nullptr;
	checkRuntime(cudaMalloc(&output_data_device, sizeof(output_data_host)));

	auto input_dims = engine->getBindingDimensions(0);
	input_dims.d[0] = input_batch;

	execution_context->setBindingDimensions(0, input_dims);
  // 用一个指针数组bindings指定input和output在gpu中的指针。
	float* bindings[] = {input_data_device, output_data_device};
	bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);

	checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream));
	checkRuntime(cudaStreamSynchronize(stream));

	float* prob = output_data_host;
	int predict_label = max_element(prob, prob + num_classes) - prob;
	float conf = prob[predict_label];
	printf("test_image: %s, max_idx: %d, probability: %f", image_path.c_str(), predict_label, conf);

  // 释放显存
	checkRuntime(cudaStreamDestroy(stream));
	checkRuntime(cudaFreeHost(input_data_host));
	checkRuntime(cudaFree(input_data_device));
	checkRuntime(cudaFree(output_data_device));
}

最终得到输出:

test_image: test.jpg, max_idx: 971, probability: 0.994527

与之前 pytorch 和 onnx 推理的结果基本一致,模型转换成功。

附录

给出完整的参考代码:https://github.com/Adenialzz/Hello-AIDeployment/tree/master/HAID/tensorrt/resnet

你可能感兴趣的:(模型部署,GPU,cuda,pytorch,c++,深度学习)