网站导航
> 文章列表 > 基于pytorch的深度学习图像分类模型部署——jetson平台

基于pytorch的深度学习图像分类模型部署——jetson平台

网友: 文章列表 2024-03-22 03:00:47

基于pytorch的深度学习图像分类模型部署——jetson平台

仓库地址:《GitHub - lxztju/pytorch_classification: 利用pytorch实现图像分类的一个完整的代码，训练，预测，TTA，模型融合，模型部署，cnn提取特征，svm或者随机森林等进行分类，模型蒸馏，一个完整的代码》

环境配置与训练测试的注意点请参考文章：《基于pytorch的深度学习图像分类》

本文将resnet18部署在jetson-nx平台，关于其他模型的部署步骤也是类似的。

resnet系列的分类模型是常用的分类模型，一般opencv-dnn也是支持的（博主没有尝试，如果大家想简单方便一点可以尝试使用最新版的opencv试下，看看是否能正常加载推理），但在jetson系列平台部署一般都是使用tensorrt进行推理。github的作者给出了两种c++的推理部署方法，分别是基于libtorch和tensorrt的。博主以tensorrt为例改写了作者的一些代码，使得阅读更为简单方便。

目录

1 模型转tensorrt引擎

1.1 第1步转onnx

1.2 onnx转tensorrt引擎

2. 平台执行tensorrt推理

1 模型转tensorrt引擎

模型转tensorrt引擎一般步骤如下：

  1. resnet.pth模型转onnx(一般用python代码转即可)
  2. 创建IBuilder的指针builder
  3. 设置推理的模式，float或者int
  4. 利用builder创建ICudaEngine的实例engine
  5. 由engine创建上下文context（即序列化）
  6. 保存context文件

其中第2-6步需要在平台端进行，以保证GPU型号的一致性。

1.1 第1步转onnx

脚本在./pytorch_classification-master\\trt_inference目录，核心代码如下

def torch_convert_onnx():model = ClsModel('resnet50', num_classes=2, dropout=0, is_pretrained=False)sd = torch.load('../cpp_inference/traced_model/trained_model.pth', map_location='cpu')model.load_state_dict(sd)model.eval()dummy_input = torch.randn(1, 3, 224, 224)torch.onnx.export(model,dummy_input,"./saved_model/torch_res50.onnx",export_params=True,input_names=["input_image"],output_names = ["model_output"],dynamic_axes = {"input_image":{0: "batch"},"model_output":{0: "batch"}},opset_version=11)

主要修改模型名称，类别数目，以及模型的路径。输入大小如果修改的话也要改过来。

1.2 onnx转tensorrt引擎

本文API以tensorrt8.x为例进行说明：

tensorrt8.x以后，默认是只支持量化到int8的，第2-6步C++核心代码如下:

int TrtClassification::build_engine(const string onnx_path, const string trt_engine_path) 
{nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger); ///创建编译引擎nvinfer1::IBuilderConfig* config = builder->createBuilderConfig(); ///创建引擎配置nvinfer1::INetworkDefinition* network=builder->createNetworkV2(1);///创建网络nvonnxparser::IParser* parser=nvonnxparser::createParser(*network, logger);///创建onnx模型解析器if (!parser->parseFromFile((char*)onnx_path.c_str(), 1)){printf("Faile to parse %s \\n", onnx_path);}config->setMaxWorkspaceSize(1 << 30);//设置每层最大可利用显存空间nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();nvinfer1::ITensor* input_tensor = network->getInput(0);nvinfer1::Dims input_dims = input_tensor->getDimensions();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);profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);config->addOptimizationProfile(profile);nvinfer1::ICudaEngine* bengine = builder->buildEngineWithConfig(*network, *config);///创建推理引擎if (bengine == nullptr){printf("build engine failed \\n");return -1;}
///将推理引擎保存下来，需要将引擎序列化，保存为二进制文件nvinfer1::IHostMemory* model_data = bengine->serialize();FILE* f = fopen((char*)trt_engine_path.c_str(), "wb");fwrite(model_data->data(), 1, model_data->size(), f);fclose(f);printf("build engine done. \\n");return 0;
}

说明：

(1). builder：构建器，搜索cuda内核目录以获得最快的可用实现，必须使用和运行时的GPU相同的GPU来构建优化引擎。在构建引擎时，TensorRT会复制权重。

nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger); ///创建编译引擎
nvinfer1::INetworkDefinition* network=builder->createNetworkV2(1);///创建网络

(2). engine：引擎，不能跨平台和TensorRT版本移植。若要存储，需要将引擎转化为一种格式，即序列化，若要推理，需要反序列化引擎。引擎用于保存网络定义和模型参数。

nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 30);
nvinfer1::ICudaEngine* bengine = builder->buildEngineWithConfig(*network, *config);
nvinfer1::IHostMemory* model_data = bengine->serialize();
FILE* f = fopen((char*)trt_engine_path.c_str(), "wb");
fwrite(model_data->data(), 1, model_data->size(), f);
fclose(f);

其中：setMaxWorkspaceSize设置比较重要，设置的小可能发挥不出并行计算的性能，设置过大可能会导致运行的显存不够而导致推理失败，甚至程序异常退出。setMaxWorkspaceSize(1 << 30)表示设置的显存大小为1G，即2的30次方，如果设置256M，则为setMaxWorkspaceSize(1 << 28)。

2. 平台执行tensorrt推理

使用推理引擎执行推理，首先需要将推理引擎反序列化，然后创建推理上下文，使用上下文执行推理过程得到推理结果，核心代码如下：

auto engine_data = load_engine_data(trt_engine_path);
auto runtime = make_nvshared(nvinfer1::createInferRuntime(logger));
engine = make_nvshared(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));///推理引擎反序列化if (engine == nullptr){printf("Deserialize cuda engine failed.\\n");return -1;}cudaStream_t stream = nullptr;checkRuntime(cudaStreamCreate(&stream));auto execution_context = make_nvshared(engine->createExecutionContext());由engine创建上下文context
  1. context：上下文，创建一些空间来存储中间值。一个engine可以创建多个context，分别执行多个推理任务。
  2. runtime：用于反序列化引擎。

需要完整推理代码的请私信我

网络标签：

相关问题