分类目录归档:pytorch

Yolov5 使用 TensorRT 进行加速

Yolov5 使用 TensorRT 进行加速

Yolov5是比较常见的机器学习模型,速度也很快,使用GPU几毫秒内就可以完成推理预测。但这就是其极限速度吗?github上Yolov5的介绍里,小模型可以跑到1ms以内,我尝试了一下,最小的模型,小图片都要3~4毫秒。后来发现要用到TensorRT速度才可以更快。下面就看看TensorRT的代码怎么写。

一.生成yolov5模型

训练生成的方法参考github上的方法进行 https://github.com/ultralytics/yolov5 。记住使用了什么大小的yolo模型和是目标识别 还是图片分类,这些后面会用到。

二.生成TensorRt可以调用的模型

生成方法参考了github上的tensorrtx项目 https://github.com/wang-xinyu/tensorrtx

1.生成wts文件

对应的python代码位置 https://github.com/wang-xinyu/tensorrtx/blob/master/yolov5/gen_wts.py

假如你是图片分类的模型,yolo训练出的模型是 best.pt,则生成wts的命令为

python gen_wts.py -w best.pt -o broken.wts -t cls

-w 输入的模型名字

-o 输出的模型名字

-t 模型类型 cls代表图片分类

2.将wts文件转换为engine文件,并导入运行。

根据我自己的理解TensorRT速度快的原因是将模型直接编译成了显卡可以直接运算的机器码,因此不同型号的显卡使用的机器码模型文件是不同的,需要特别生成。

github上的tensorrtx项目虽然有c++的实现代码,但没有采用c++类的写法,而且有许多公共变量,不方便同一个程序,调用不同类型 不同大小的模型文件。因此进行了部分改写。文件结构如下图

全部文件打包在如下压缩包里面

生成模型的使用方法如下:

#include "yolov5det.h"
#include "yolov5cls.h"

int main(int argc, char** argv) {
    yolov5det _yolov5det;
    yolov5cls _yolov5cls;

    string outputfilename = "best.engine";//输出
    string engine = std::string("best.wts");//输入
    _yolov5det.createngine(engine,outputfilename,'s');// s 表示 yolov5s 模型

    string outputfilename = "cls.engine";//输出
    string engine = std::string("cls.wts");//输入
    _yolov5cls.createngine(engine,outputfilename,'s');// s 表示 yolov5s 模型
}

进行预测的使用方法

#include "yolov5det.h"
#include "yolov5cls.h"

#include "opencv2/imgcodecs.hpp"

using namespace std;
using namespace cv;

int main(int argc, char** argv) {
    yolov5det _yolov5det;
    yolov5cls _yolov5cls;
    clock_t starttime, endtime;

    string enginedet = "best.engine";
    _yolov5det.setengine(enginedet);//设置模型
    _yolov5det.kConfThresh = 0.75;//设置kConfThresh参数
    _yolov5det.kNmsThresh = 0.35;//设置kNmsThresh参数
	
	
    Mat img_det = imread("det.bmp",IMREAD_COLOR); //导入图片
    starttime = clock();
    std::vector<std::vector<Detection>> res_batch = _yolov5det.detect(img_det); //预测
    endtime = clock();
    cout << (double)(endtime - starttime) << endl;

    if(res_batch.size()>0){
        for(int n =0;n<res_batch[0].size();n++){//res_batch[0] 测试的模型输入只有1张图片
            cout<<res_batch[0][n].bbox[3]<<endl;//res_batch[0][n] 的属性里有box的位置,分类,分值等信息
            cout<<res_batch[0][n].bbox[2]<<endl;
        }
    }


    string enginecls = "cls.engine";
    _yolov5cls.setengine(enginecls);//设置模型
	
    Mat img_cls = imread("cls.bmp",IMREAD_COLOR);
    starttime = clock();
    std::vector<std::vector<int>> res_batch = _yolov5cls.detect(img_cls);
    endtime = clock();
    cout << (double)(endtime - starttime) << endl;

    if(res_batch.size()>0){
        cout<<"classid:"<<res_batch[0][0]<<endl;// 分类的id
    }
}
使用TorchSharp 调用 PyTorch 语义分割神经网络模型

使用TorchSharp 调用 PyTorch 语义分割神经网络模型

TorchSharp 是对 Torch c++的封装,基本继承了c++的全部接口。但使用中会有一些小问题,需要特别注意一些。

  1. 语义分割(semantic segmentation)神经网络训练

训练的代码可以参考github里的官方代码 https://github.com/pytorch/vision/tree/main/references/segmentation

2.模型输出

官方代码的模型 默认输出是list 虽然可以强制输出script文件,但TorchSharp 调用后会报错”Expected Tensor but got GenericDict”.因此需要修改网络

export代码如下:

import torch
import torchvision
from torch import nn
import numpy as np

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self._model = torchvision.models.segmentation.lraspp_mobilenet_v3_large(num_classes=2,
                aux_loss=False,
                pretrained=False)
        
        #修改了输入,将输入改为单通道图片
        #如果输入的是3通道 则不需要修改
        for item in self._model.backbone.items():          
            item[1][0] = nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
            break                                                         
                                                                 
        checkpoint = torch.load('model/model_119.pth', map_location='cpu')
        self._model.load_state_dict(checkpoint['model'], strict=not True)
        self._model.eval()
        
    def forward(self, x):
        # 修改了输出,将List修改为Tensor 并进行了 argmax 和 转float操作
        result = self._model.forward(x)
        return result["out"].argmax(1).flatten().float()


model = MyModel()
x = torch.rand(1,1, 240, 400)
predictions = model(x)

#使用torch.jit.trace 输出 script pt 网络文件
traced_script_module = torch.jit.trace(model, x)
traced_script_module.save('_seg.pt') 

3.使用TorchSharp 进行预测

由于测试程序使用了opencvsharp,所以下面的代码使用了opencv来读取图片和简单的数据预处理,并使用opencv可视化输出

//导入网络
torch.jit.ScriptModule torch_model;
torch_model = torch.jit.load("_seg.pt");

//导入图片,并使用BlobFromImage 进行数据类型 维度 和归一化的转换
Mat temp = Cv2.ImRead("Pic_42633.bmp", ImreadModes.Grayscale);
Mat tensor_mat = OpenCvSharp.Dnn.CvDnn.BlobFromImage(temp, 1 / 255.0);

//初始化输入的tensor
float[] data_bytes = new float[tensor_mat.Total()];
Marshal.Copy(tensor_mat.Data, data_bytes, 0, data_bytes.Length);
torch.Tensor x = torch.tensor(data_bytes, torch.ScalarType.Float32);

//维度转换 和 normalize(进行normalize是因为官方代码里有这一步处理)
x = x.reshape(1, 1, 240, 400);
x = TorchSharp.torchvision.transforms.functional.normalize(x, new double[] { 0.485 }, new double[] { 0.229 });

DateTime date1 = DateTime.Now;
//进行预测
torch.Tensor _out = torch_model.forward(x);
DateTime date2 = DateTime.Now;
TimeSpan ts = date2 - date1;
Console.WriteLine("No. of Seconds (Difference) = {0}", ts.TotalMilliseconds);
Console.WriteLine(_out);

//使用opencv 将预测出的tensor输出为可视化的图片
Mat result_mat = new Mat(240, 400, MatType.CV_32FC1, _out.bytes.ToArray());
result_mat.ConvertTo(result_mat, MatType.CV_8UC1);
Cv2.ImWrite("mask.bmp", result_mat);