篇章五:TensorRT的应用场景及部署模块应用
目录
PS:纯粹为学习分享经验,不参与商用价值运作,若有侵权请及时联系!!!
深度学习模型部署TensorRT加速(六):TensorRT部署自定义CNN模型
深度学习模型部署TensorRT加速(七):TensorRT部署一个图像分类模型
深度学习模型部署TensorRT加速(八):TensorRT部署目标检测YOLO模型
一、TensorRT 部署流程:
回顾:深度学习模型部署TensorRT加速(三): TensorRT模型部署及优化
·TensorRT 部署流程主要有以下五步:
1.训练模型
2.导出模型为 ONNX 格式
3.选择精度
4.转化成 TensorRT 模型
5.部署模型
训练模型:一般训练可以在基础深度学习模型中完成,获取其权重文件。
下文着重介绍导出模型为 ONNX 格式和转化成 TensorRT 模型的执行流程!!!!!!
导出模型为 ONNX 格式:
模型部署的常见流水线是“深度学习框架-中间表示-推理引擎”。其中比较常用的一个中间表示是 ONNX。
前期:创建 PyTorch 模型
中期:中间表示 - ONNX
下面是一个简单的示例代码,演示如何将PyTorch模型导出为ONNX模型:
import torch
import torchvision
# 加载并训练PyTorch模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()
# 创建一个示例输入张量
dummy_input = torch.randn(1, 3, 224, 224)
# 导出模型为ONNX格式
onnx_path = "model.onnx"
torch.onnx.export(model, dummy_input, onnx_path, export_params=True, opset_version=11)
print("ONNX模型已导出:", onnx_path)
PyTorch 框架自带对 ONNX 的支持,只需要构造一组随机的输入,并对模型调用 torch.onnx.export 即可完成 PyTorch 到 ONNX 的转换。
推理引擎 ONNX Runtime 对 ONNX 模型有原生的支持。给定一个 .onnx 文件,只需要简单使用 ONNX Runtime 的 Python API 就可以完成模型推理。
ONNX模型可视化网站:Netron
输入ONNX文件可以得到其结构图及对应输出:
后期:ONNX模型转化为TensorRT模型
可以使用TensorRT的Python API或C++ API
Python代码示例:
1.加载ONNX模型
import tensorrt as trt
# 加载ONNX模型
onnx_path = "model.onnx"
onnx_model = onnx.load(onnx_path)
2.创建TensorRT的Builder对象和NetworkDefinition对象
# 创建TensorRT的Builder对象
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
# 创建TensorRT的NetworkDefinition对象
network = builder.create_network()
3.解析ONNX模型: 使用TensorRT的解析器(Parser)将ONNX模型转换为TensorRT的网络。
# 创建TensorRT的ONNX解析器
parser = trt.OnnxParser(network, trt.Logger(trt.Logger.WARNING))
# 解析ONNX模型并将其转换为TensorRT网络
parser.parse(onnx_model.SerializeToString())
4.构建TensorRT的Engine: 使用TensorRT的Builder对象构建TensorRT的Engine。
# 构建TensorRT的Engine
engine = builder.build_cuda_engine(network)
5.保存TensorRT模型
# 保存TensorRT模型到文件
trt_path = "model.trt"
with open(trt_path, "wb") as f:
f.write(engine.serialize())
完成上述步骤后,将获得一个转换为TensorRT格式的模型文件(model.trt)。可以将该文件用于TensorRT的推理和部署。
同时补上C++ 将ONNX模型转换为TensorRT模型的示例代码:
#include <iostream>
#include <fstream>
#include <NvInfer.h>
#include <onnx/onnx.pb.h>
#include <onnxparser/OnnxParser.h>
int main()
{
// 读取ONNX模型文件
std::string onnxPath = "model.onnx";
std::ifstream onnxFile(onnxPath, std::ios::binary);
if (!onnxFile)
{
std::cerr << "无法打开ONNX模型文件: " << onnxPath << std::endl;
return 1;
}
onnx::ModelProto onnxModel;
if (!onnxModel.ParseFromIstream(&onnxFile))
{
std::cerr << "无法解析ONNX模型文件: " << onnxPath << std::endl;
return 1;
}
// 创建TensorRT的Builder对象
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
// 创建TensorRT的NetworkDefinition对象
nvinfer1::INetworkDefinition* network = builder->createNetwork();
// 创建TensorRT的ONNX解析器
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
// 解析ONNX模型并将其转换为TensorRT网络
if (!parser->parse(onnxModel))
{
std::cerr << "无法解析ONNX模型" << std::endl;
return 1;
}
// 构建TensorRT的Engine
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// 保存TensorRT模型到文件
std::string trtPath = "model.trt";
std::ofstream trtFile(trtPath, std::ios::binary);
if (!trtFile)
{
std::cerr << "无法打开TensorRT模型文件: " << trtPath << std::endl;
return 1;
}
nvinfer1::IHostMemory* modelStream = engine->serialize();
trtFile.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
trtFile.close();
// 释放资源
modelStream->destroy();
parser->destroy();
network->destroy();
config->destroy();
engine->destroy();
builder->destroy();
std::cout << "TensorRT模型已保存: " << trtPath << std::endl;
return 0;
}
经过以上的学习,目前已经知道TensorRT是如何解析权重文件了,接下来将注意展示具体的深度学习模型实操环节!!!!从简单的模型部署到实际应用,Action!!!!!
二、构建ResNet 模型部署及推理
构建、部署和执行ResNet模型的推理通常需要几个步骤:模型构建、模型训练、模型导出、模型部署和推理。以下是一个简单的示例,演示了如何使用Python和PYTORCH构建、导出、部署和执行ResNet模型的推理。
1.模型构建和训练:
import torch
import torch.nn as nn
import torchvision.models as models
# 构建 ResNet50 模型
model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(num_ftrs, 1024),
nn.ReLU(),
nn.Linear(1024, num_classes)
)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练模型
# ...
2.模型导出:
# 保存训练好的模型
torch.save(model.state_dict(), "resnet_model.pth")
模型部署及推理:
import torch
import cv2
import numpy as np
# 加载部署的模型
model = YourDeployedModel() # 在这里加载您的模型
# 准备输入数据
input_image = cv2.imread("path/to/input_image.jpg")
input_image = cv2.resize(input_image, (224, 224)) # 调整输入图像大小
input_image = np.transpose(input_image, (2, 0, 1)) # 转换为CHW格式
input_image = torch.from_numpy(input_image).unsqueeze(0).float() # 添加批次维度
# 执行推理
with torch.no_grad():
output = model(input_image)
# 处理推理结果
# ...
在安装好tensorrt环境后,可以尝试使用预训练权重进行转化封装部署,运行以下代码!!
具体完整示例:ResNet18
Python代码
import torch
import torch.nn as nn
import torchvision.models as models
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
# Step 1: 加载模型
resnet = models.resnet18(pretrained=True)
resnet.eval()
# Step 2: 导出为ONNX格式
input_shape = (1, 3, 224, 224)
input_data = torch.randn(input_shape)
torch.onnx.export(resnet, input_data, "resnet18.onnx", verbose=True)
# Step 3: 构建TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
with open("resnet18.onnx", "rb") as model:
parser.parse(model.read())
builder.max_workspace_size = 1 << 30
engine = builder.build_cuda_engine(network)
# Step 4: 推理
input_shape = (1, 3, 224, 224)
input_data = torch.randn(input_shape).cuda().numpy()
output_shape = (1, 1000)
output_data = np.empty(output_shape, dtype=np.float32)
with engine.create_execution_context() as context:
bindings = [int(input_data.ctypes.data), int(output_data.ctypes.data)]
context.execute_v2(bindings)
print("Inference Result:", output_data.argmax())
PS:纯粹为学习分享经验,不参与商用价值运作,若有侵权请及时联系!!!
下篇内容预告:
-
深度学习模型部署TensorRT加速(六):TensorRT部署自定义CNN模型
-
深度学习模型部署TensorRT加速(七):TensorRT部署一个图像分类模型
-
深度学习模型部署TensorRT加速(八):TensorRT部署目标检测YOLO模型