【TRT】使用TensorRT进行分类模型推理
1. pytorch 模型导出为onnx模型
1.1 pytorch 模型代码
import torch
import torchvision
import cv2
import numpy as np
class Classifier(torch.nn.Module):
def __init__(self):
super().__init__()
#使用torchvision自带的与训练模型, 更多模型请参考:https://tensorvision.readthedocs.io/en/master/
self.backbone = torchvision.models.resnet18(pretrained=True)
def forward(self, x):
feature = self.backbone(x)
// 将softmax 加入到模型，省去推理时后处理的归一化操作
probability = torch.softmax(feature, dim=1)
return probability
- 将softmax 加入到模型，省去推理时后处理的归一化操作
1.2 导出onnx模型
dummy = torch.zeros(1, 3, 224, 224)
torch.onnx.export(
model, (dummy,), "workspace/classifier.onnx",
input_names=["image"],
output_names=["prob"],
dynamic_axes={"image": {0: "batch"}, "prob": {0: "batch"}},
opset_version=11
)
- dynamic_axes 只指定 batch
1.3 模型推理代码
# 对每个通道进行归一化有助于模型的训练
imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std = [0.229, 0.224, 0.225]
image = cv2.imread("workspace/dog.jpg")
image = cv2.resize(image, (224, 224)) # resize
image = image[..., ::-1] # BGR -> RGB
image = image / 255.0
image = (image - imagenet_mean) / imagenet_std # normalize
image = image.astype(np.float32) # float64 -> float32
image = image.transpose(2, 0, 1) # HWC -> CHW
image = np.ascontiguousarray(image) # contiguous array memory
image = image[None, ...] # CHW -> 1CHW
image = torch.from_numpy(image) # numpy -> torch
model = Classifier().eval()
with torch.no_grad():
probability = model(image)
predict_class = probability.argmax(dim=1).item()
confidence = probability[0, predict_class]
labels = open("workspace/labels.imagenet.txt", encoding='utf-8').readlines()
labels = [item.strip() for item in labels]
print(f"Predict: {predict_class}, {confidence}, {labels[predict_class]}")
推理结果
Predict: 263, 0.32459262013435364, 彭布洛克威尔士科基犬
2. C++ TensorRT推理代码
2.1 前处理代码
- resize
- 归一化
- bgr 转 rgb 转 tensor （RRR GGG BBB）
/* python 代码
image = cv2.resize(image, (224, 224)) # resize
image = image[..., ::-1] # BGR -> RGB
image = image / 255.0
image = (image - imagenet_mean) / imagenet_std # normalize
image = image.astype(np.float32) # float64 -> float32
image = image.transpose(2, 0, 1) # HWC -> CHW
image = np.ascontiguousarray(image) # contiguous array memory
image = image[None, ...] # CHW -> 1CHW
*/
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)));
// 归一化，通道转换， toTensor
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;
// opencv 的存储格式为 BGR BGR BGR BGR
// 转换为 tensor 的 [BBB] [GGG] [RRR]
// 在此过程中再进行图像的归一化, bgr 转 rrr ggg bbb
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));
3. 推理代码
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 = execution_context->getBindingDimensions(0);
input_dims.d[0] = input_batch;
// 设置当前推理时，input大小
execution_context->setBindingDimensions(0, input_dims);
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 = std::max_element(prob, prob + num_classes) - prob; // 确定预测类别的下标
auto labels = load_labels("labels.imagenet.txt");
auto predict_name = labels[predict_label];
float confidence = prob[predict_label]; // 获得预测值的置信度
printf("Predict: %s, confidence = %f, label = %d\n", predict_name.c_str(), confidence, predict_label);
checkRuntime(cudaStreamDestroy(stream));
checkRuntime(cudaFreeHost(input_data_host));
checkRuntime(cudaFree(input_data_device));
checkRuntime(cudaFree(output_data_device));
- 申请 输入的host memory 和 global memory
- 在CPU上处理完成或，拷贝到 GPU
- 设置input_dim 的输入尺寸
- 绑定输入输出的GPU内存
- 将GPU结果拷贝到 host
- CPU 后处理
- 释放内存
完整代码
// tensorRT include
// 编译用的头文件
#include