南昌的网站设计,做网站精英,合肥网站建设发布,wordpress制作左侧边栏目录
简介
分析PyTorch示例
onnx模型转engine
编写TensorRT推理代码
main.cpp测试代码
小结 简介
这里通过TensorRTUNet#xff0c;在Linux下实现对遥感图像的变化检测#xff0c;示例如下#xff1a; 可以先拉去代码#xff1a;RemoteChangeDetection 分析PyTorch示…目录
简介
分析PyTorch示例
onnx模型转engine
编写TensorRT推理代码
main.cpp测试代码
小结 简介
这里通过TensorRTUNet在Linux下实现对遥感图像的变化检测示例如下 可以先拉去代码RemoteChangeDetection 分析PyTorch示例
在目录PyFiles中unet.py存放UNet网络定义可以使用test_infer.py脚本进行推理并导出onnx模型可以简单分析一下test_infer.py中的关键代码。
1加载处理图像
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from unet import UNet
import onnx
import onnxsim# 读取变换前后的代码
img1 Image.open(./A/val_20.png)
img2 Image.open(./B/val_20.png)# 输出的图像名称
output_image_path result.png# PIL图像转Tensor张量
transform transforms.Compose([transforms.ToTensor()
])# 分别取两幅图像的第一个通道图像因为PIL读取的图像是RGB的注意和OpenCV图像区别
img1_data np.array(img1)
img1_data img1_data[:, :, 0]img2_data np.array(img2)
img2_data img2_data[:, :, 0]# 这里合并输入图像: shape [height, width, 2]
input_image np.stack([img1_data, img2_data], axis2)# 转换为模型输入大致流程:
# 1. transform: 图像从[0, 255] 映射到 [0, 1]; 交换通道图像[h, w, 2] [2, h, w]
# 2. unsqueeze(0)增加第一个维度[2, h, w] [1, 2, h, w]
# 3. unit8 转 float32类型并放置在GPU上
input_image_tensor transform(input_image).unsqueeze(0).type(torch.float32).to(device)
2推理并导出为onnx
def export_norm_onnx(model, file, input):torch.onnx.export(model model, args (input,),f file,input_names [input0],output_names [output0],opset_version 9)print(Finished normal onnx export)model_onnx onnx.load(file)onnx.checker.check_model(model_onnx)# 使用onnx-simplifier来进行onnx的简化。print(fSimplifying with onnx-simplifier {onnxsim.__version__}...)model_onnx, check onnxsim.simplify(model_onnx)assert check, assert check failedonnx.save(model_onnx, file)
这里定义了一个导出onnx函数model为PyTorch模型file是输出文件路径input是模型的输入。
with torch.no_grad():net UNet(2).to(device)net.eval()load_models torch.load(weights)net.load_state_dict(torch.load(weights))out_image net(input_image_tensor)_out_image out_image[0][0].round().detach().cpu().numpy()_out_image (_out_image * 255).astype(np.uint8)result_image Image.fromarray(_out_image)result_image.save(output_image_path)export_norm_onnx(net, ./unet_simple.onnx, input_image_tensor)
这里是推理为了测试.pth模型并导出onnx。这里注意对输出图像的后处理过程在编写c接口时要留意。 使用onnx可视化工具查看导出的onnx模型 onnx模型转engine
如果你已经按照了TensorRT并且配置好了环境变量后可以直接使用bin下的trtexec命令将onnx模型进行转换假如你的TensorRT安装路径如下 环境变量的配置 使用如下命令进行转换
trtexec --onnxdncnn_color_blind.onnx --saveEnginedncnn_color_engine_intro.engine --explicitBatch
// *.onnx是输入的模型*.engine是保存的模型
上边只是举个例子把文件名换成自己的就可以了。
编写TensorRT推理代码
1运行环境搭建
我的运行环境目录大致如下
RemoteChangeDetection
3rdparty|------- opencv-3.4.10|-------- include|-------- lib|------- TensorRT-8.5.2.2|-------- include|-------- lib...
首先修改CMakeLists.txt中的三方库路径 那么你应该修改CUDACUDNNOpenCV以及TensorRT的路径。
在src/路径下是核心代码trt_logger包含了TensorRT推理时依赖的logger以及CUDA函数运行时的检查宏
#ifndef __LOGGER_H__
#define __LOGGER_H__#include string
#include stdarg.h
#include memory
#include cuda_runtime.h
#include system_error
#include NvInfer.h#define CUDA_CHECK(call) __cudaCheck(call, __FILE__, __LINE__)
#define LAST_KERNEL_CHECK(call) __kernelCheck(__FILE__, __LINE__)static void __cudaCheck(cudaError_t err, const char* file, const int line) {if (err ! cudaSuccess) {printf(ERROR: %s:%d, , file, line);printf(code:%s, reason:%s\n, cudaGetErrorName(err), cudaGetErrorString(err));exit(1);}
}static void __kernelCheck(const char* file, const int line) {cudaError_t err cudaPeekAtLastError();if (err ! cudaSuccess) {printf(ERROR: %s:%d, , file, line);printf(code:%s, reason:%s\n, cudaGetErrorName(err), cudaGetErrorString(err));exit(1);}
}#define LOGF(...) logger::Logger::__log_info(logger::Level::FATAL, __VA_ARGS__)
#define LOGE(...) logger::Logger::__log_info(logger::Level::ERROR, __VA_ARGS__)
#define LOGW(...) logger::Logger::__log_info(logger::Level::WARN, __VA_ARGS__)
#define LOG(...) logger::Logger::__log_info(logger::Level::INFO, __VA_ARGS__)
#define LOGV(...) logger::Logger::__log_info(logger::Level::VERB, __VA_ARGS__)
#define LOGD(...) logger::Logger::__log_info(logger::Level::DEBUG, __VA_ARGS__)
最重要的是UNetTrt部分在UNetTrt.h
#ifndef UNET_TRT_H_
#define UNET_TRT_H_#include iostream
#include memory
#include opencv2/opencv.hpp
#include cuda_runtime.h// 前置定义
namespace nvinfer1
{class IRuntime;class ICudaEngine;class IExecutionContext;
}class UNet
{
public:UNet() {};~UNet();// 加载engine文件bool loadTrtModel(const std::string model_path);// 推理input_mat1: 变换前input_mat2: 变换后output是变量引用bool trt_infer(cv::Mat input_mat1, cv::Mat input_mat2, cv::Mat output); // input_mat1: before, input_mat2: afterprivate:// runtime_, engine_, context_等成员是TensorRT推理时最重要的几个成员变量// 为了放置内存泄露用智能指针管理std::shared_ptrnvinfer1::IRuntime runtime_; std::shared_ptrnvinfer1::ICudaEngine engine_;std::shared_ptrnvinfer1::IExecutionContext context_;cudaStream_t stream_;int input_index_; // 索引输入int output_index_; // 索引输出const char *INPUT_NAME input0; // 输入名称和onnx导入时保持一致const char *OUTPUT_NAME output0; // 和上边保持一致const int BATCH_SIZE 1; // 一般都保持为1void *buffers_[2]; // 存放TensorRT输入输出float *input_float_ nullptr; // 存放Host端输入c11允许.h中初始化float *output_float_ nullptr; // Host端计算结果
};#endif
在.cpp中给出一些核心实现
#include UNetTrt.h
#include fstream
#include cmath
#include trt_logger.h
#include NvInfer.h
#include NvOnnxParser.h#define INPUT_WIDTH 1024
#define INPUT_HEIGHT 1024bool UNet::loadTrtModel(const std::string model_path)
{char *trt_stream nullptr;size_t size 0;// load trt modelstd::ifstream file(model_path, std::ios::binary);if (file.good()) {file.seekg(0, file.end);size file.tellg();file.seekg(0, file.beg);trt_stream new char[size];if(!trt_stream)return false;file.read(trt_stream, size);file.close();} else {return false;}logger::Logger trt_logger(logger::Level::INFO);runtime_.reset(nvinfer1::createInferRuntime(trt_logger));if(!runtime_)return false;engine_.reset(runtime_-deserializeCudaEngine(trt_stream, size, nullptr));if(!engine_)return false;context_.reset(engine_-createExecutionContext());if(!context_)return false;const nvinfer1::ICudaEngine trtEngine context_-getEngine();input_index_ trtEngine.getBindingIndex(INPUT_NAME);output_index_ trtEngine.getBindingIndex(OUTPUT_NAME);CUDA_CHECK(cudaMalloc(buffers_[input_index_], BATCH_SIZE * 2 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float)));CUDA_CHECK(cudaMalloc(buffers_[output_index_], BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float)));input_float_ new float[BATCH_SIZE * 2 * INPUT_WIDTH * INPUT_HEIGHT];output_float_ new float[BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT];delete []trt_stream;return true;
}
首先输入大小是固定的所以在宏里写死了输入大小1024x1024loadTrtModel根据路径加载engine文件并对一些推理时用到的成员变量依次初始化同时分配好输入输出空间。
推理代码如下
bool UNet::trt_infer(cv::Mat input_mat1, cv::Mat input_mat2, cv::Mat output)
{if(input_mat1.empty() || input_mat2.empty())return false;if(input_mat1.rows ! input_mat2.rows || input_mat1.cols ! input_mat2.cols)return false;if(input_mat1.channels() 1 input_mat2.channels() 1) return false;int pre_width input_mat1.cols;int pre_height input_mat1.rows;cv::resize(input_mat1, input_mat1, cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::INTER_CUBIC);cv::resize(input_mat2, input_mat2, cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::INTER_CUBIC);std::vectorcv::Mat input_mat1_channels;cv::split(input_mat1, input_mat1_channels);std::vectorcv::Mat input_mat2_channels;cv::split(input_mat2, input_mat2_channels);// [H, W, C] [C, H, W] [0.0, 0.1]for(int i 0; i INPUT_WIDTH; i) {for(int j 0; j INPUT_HEIGHT; j) {int idx_c1 j * INPUT_WIDTH i;int idx_c2 idx_c1 INPUT_WIDTH * INPUT_HEIGHT;input_float_[idx_c1] (float)input_mat1_channels[2].data[idx_c1] / 255.0f;input_float_[idx_c2] (float)input_mat2_channels[2].data[idx_c1] / 255.0f;}}memset(output_float_, 0, BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT);CUDA_CHECK(cudaStreamCreate(stream_));CUDA_CHECK(cudaMemcpyAsync(buffers_[input_index_], input_float_, BATCH_SIZE * 2 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float), cudaMemcpyHostToDevice, stream_));context_-enqueueV2(buffers_, stream_, nullptr);CUDA_CHECK(cudaMemcpyAsync(output_float_, buffers_[output_index_], BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float), cudaMemcpyDeviceToHost, stream_));cudaStreamSynchronize(stream_);// roundfor(int i 0; i INPUT_WIDTH; i) {for(int j 0; j INPUT_HEIGHT; j) {int index j * INPUT_WIDTH i;output_float_[index] std::round(output_float_[index]);}}output cv::Mat(INPUT_HEIGHT, INPUT_WIDTH, CV_32F, output_float_);output * 255.0;output.convertTo(output, CV_8U);cv::resize(output, output, cv::Size(pre_width, pre_height), cv::INTER_CUBIC);return true;
}
这里依次讲解一下首先你可能要把代码放入工程那么应该尽量做好判断比如图像是否为空图像大小、通道是否一致以防万一可以同时进行Resize
cv::split对3通道图像进行剥离放入vector中然后开始进行通道转换与归一化。这里可以稍微理解一下不同图像在内存中的存放方式一般的RGB图像或者BGR图像height, width, channel应该是这样
B G R B G R B G R B G R B G R B G R
B G R B G R B G R B G R B G R B G R
B G R B G R B G R B G R B G R B G R
B G R B G R B G R B G R B G R B G R互相交错存放但是网络输入一般是channel, height, width那么存放方式是如下这样
R R R R R R R R
R R R R R R R R
R R R R R R R RG G G G G G G G
G G G G G G G G
G G G G G G G GB B B B B B B B
B B B B B B B B
B B B B B B B B那么就可以很容易写出通道转换与归一化代码
// [H, W, C] [C, H, W] [0.0, 0.1]for(int i 0; i INPUT_WIDTH; i) {for(int j 0; j INPUT_HEIGHT; j) {int idx_c1 j * INPUT_WIDTH i;int idx_c2 idx_c1 INPUT_WIDTH * INPUT_HEIGHT;input_float_[idx_c1] (float)input_mat1_channels[2].data[idx_c1] / 255.0f;input_float_[idx_c2] (float)input_mat2_channels[2].data[idx_c1] / 255.0f;}}
每次推理前把输出结果清空置为0
memset(output_float_, 0, BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT);
重新分配cudaStream_tcudaMemcpyAsync分配显存context_-enqueueV2推理cudaMemcpyAsync再将结果从显存拷贝到Host端。
CUDA_CHECK(cudaStreamCreate(stream_));CUDA_CHECK(cudaMemcpyAsync(buffers_[input_index_], input_float_, BATCH_SIZE * 2 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float), cudaMemcpyHostToDevice, stream_));context_-enqueueV2(buffers_, stream_, nullptr);CUDA_CHECK(cudaMemcpyAsync(output_float_, buffers_[output_index_], BATCH_SIZE * 1 * INPUT_WIDTH * INPUT_HEIGHT * sizeof(float), cudaMemcpyDeviceToHost, stream_));cudaStreamSynchronize(stream_);
后处理过程中也遇到了一些坑总体而言还是要一一对照python那部分后处理代码仔细分析
_out_image out_image[0][0].round().detach().cpu().numpy()
_out_image (_out_image * 255).astype(np.uint8)
result_image Image.fromarray(_out_image)
result_image.save(output_image_path) // roundfor(int i 0; i INPUT_WIDTH; i) {for(int j 0; j INPUT_HEIGHT; j) {int index j * INPUT_WIDTH i;output_float_[index] std::round(output_float_[index]);}}output cv::Mat(INPUT_HEIGHT, INPUT_WIDTH, CV_32F, output_float_);output * 255.0;output.convertTo(output, CV_8U);cv::resize(output, output, cv::Size(pre_width, pre_height), cv::INTER_CUBIC);
因为图像是缩放过一次的最后给缩放回去。
main.cpp测试代码
在main.cpp编写测试示例一般是建议将类用智能指针管理
std::shared_ptrUNet unet_infer std::make_sharedUNet();std::string model_path ./weights/unet_simple_trt.engine;if(unet_infer) {if(unet_infer-loadTrtModel(model_path))std::cout UNet Init Successful! \n;else std::cout UNet Init Failed! \n;
}
推理
cv::Mat img1 cv::imread(./test_images/val_20_A.png);
cv::Mat img2 cv::imread(./test_images/val_20_B.png);
cv::Mat result;if(unet_infer-trt_infer(img1, img2, result)) {std::cout UNet Infer Successfully! \n;
} else {std::cout UNet Infer Failed! \n;
}
当然最后可以测试一下推理速度以及输出是不是一致
int count 100;
int cost 0;for(int i 0; i count; i) { auto start std::chrono::high_resolution_clock::now(); bool success unet_infer-trt_infer(img1, img2, result);auto end std::chrono::high_resolution_clock::now();cost std::chrono::duration_caststd::chrono::milliseconds(end - start).count();
}std::cout duration: (float)(cost) / count ms std::endl; if(!result.empty()) {cv::imwrite(./result.png, result);
} 显存占用 对比了一下1024x1024的输入大概会消耗1G的显存如果你缩小图像后再计算效果会差一些。
计算耗时 大概是80~90ms一张吧。
小结
上边只是初步实现了变换检测的推理但是图像预处理与后处理还是有很多可以优化改进的地方后边有时间再补上吧。
参考资料
ESCNet
ChangeDetection_GUI
TensorRT
tensorrtx
基于CUDA的并行计算技术
