使用TensotRT 4步搞定mobilnetV3深度学习模型部署

前言

MobilnetV3作为谷歌目前发布的最新的轻量级深度学习模型，在性能和模型大小上都达到了很好的平衡，从而使其在移动端部署具有很大的优势，这里我们使用C++ 和Cuda模型推理加速工具TensorRT将训练好的Mobilenet模型封装成dll，使其能够方便快捷的部署在移动端。整过过程可以分为以下几个步骤：

• step 1 配置环境，安装TensorRT
• step 2 训练模型，并将模型转为.engine格式
• step 3 将Mobilenet封装成一个动态链接库dll
• step 4 测试

1. 配置环境，安装TensorRT

假设你已经了解了Pytorch基础知识，并配置好了MobilenetV3源码所需要的环境，与此同时还配置了TensorRT的环境并测试其安装成功。如果这些还没有掌握可以参考以下博文：
1，深度学习之Pytorch环境搭建
2，TensorRT安装指南
3，yolov5部署之七步完成tensorRT模型推理加速

2. 训练模型，并将模型转为.engine格式

关于源码的选择，github上有很多选择，这里选用了官方的源码，但只有模型文件，因此train文件需要结合自己的喜好进行编写。

• train.py

import os
import time
import argparse
import logging
import numpy as np
from collections import OrderedDict
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from ghostnet import ghostnet
from sklearn.metrics import accuracy_score
from mobilenetv3 import mobilenetv3_small


def accuracy(y_pred,y_true):
    y_pred_cls = torch.argmax(nn.Softmax(dim=1)(y_pred),dim=1).data
    return accuracy_score(y_true,y_pred_cls)

def train(valdir):
    # normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
    #                                  std=[0.229, 0.224, 0.225])
    loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, transforms.Compose([
            # transforms.Resize(160),
            # transforms.CenterCrop(160),
            transforms.ToTensor(),
            # normalize,
        ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    # model = ghostnet(num_classes=args.num_classes, width=args.width, dropout=args.dropout)
    model = mobilenetv3_small(num_classes=args.num_classes)
    model.optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    model.loss_func = nn.CrossEntropyLoss()
    # model.metric_func = accuracy
    # model.metric_name = "accuracy"
    epochs = 10
    log_step_freq = 100

    for epoch in range(1, epochs + 1):
        # 1，训练循环-------------------------------------------------
        model.train()
        loss_sum = 0.0
        metric_sum = 0.0
        step = 1
        for step, (features, labels) in enumerate(loader, 1):
            # 梯度清零
            model.optimizer.zero_grad()

            # 正向传播求损失
            predictions = model(features)
            loss = model.loss_func(predictions, labels)
            metric = accuracy(predictions, labels)

            # 反向传播求梯度
            loss.backward()
            model.optimizer.step()

            # 打印batch级别日志
            loss_sum += loss.item()
            metric_sum += metric.item()
            if step % log_step_freq == 0:
                print(("[step = %d] loss: %.3f, " +"accruacy" + ": %.3f") %
                      (step, loss_sum / step, metric_sum / step))
    torch.save(model, './models/model.pt')


def valid_step(model, features, labels):
    # 预测模式，dropout层不发生作用
    model.eval()

    predictions = model(features)
    loss = model.loss_func(predictions, labels)
    metric = model.metric_func(predictions, labels)

    return loss.item(), metric.item()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch ImageNet Inference')
    parser.add_argument('--data', metavar='DIR', default='./data',
                        help='path to dataset')
    parser.add_argument('--output_dir', metavar='DIR', default='./models',
                        help='path to output files')
    parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
                        help='number of data loading workers (default: 2)')
    parser.add_argument('-b', '--batch-size', default=8, type=int,
                        metavar='N', help='mini-batch size (default: 1)')
    parser.add_argument('--num-classes', type=int, default=2,
                        help='Number classes in dataset')
    parser.add_argument('--width', type=float, default=1.0,
                        help='Width ratio (default: 1.0)')
    # parser.add_argument('--dropout', type=float, default=0.2, metavar='PCT',
    #                     help='Dropout rate (default: 0.2)')
    parser.add_argument('--num-gpu', type=int, default=1,
                        help='Number of GPUS to use')
    args = parser.parse_args()

    valdir = os.path.join(args.data, 'train')
    train(valdir)

这里使用的是Pytorch框架，训练完成后将模型保存为.pt格式，然后将.pt模型转为.wts模型，最后将.wts模型转为.engine模型部署在移动端，.pt转.wts模型的代码如下：

• gen_wts.py

import torch
import struct
from utils.torch_utils import select_device

# Initialize
device = select_device('cpu')
# Load model
model = torch.load('models/model.pt', map_location=device) # load to FP32
model.to(device).eval()

f = open('model.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
    vr = v.reshape(-1).cpu().numpy()
    f.write('{} {} '.format(k, len(vr)))
    for vv in vr:
        f.write(' ')
        f.write(struct.pack('>f',float(vv)).hex())
    f.write('\n')

3. 将Mobilenet封装成一个动态链接库dll

首先新建一个dll项目，然后再新建一个控制台应用来调用dll，假设你已经了解VS2015建立不同类型项目的基础知识以及熟悉常规的类封装方法，如果这些基础知识还不具备可参考以下博文：
1，私有类封装为DLL的方法
2，抽象类作为接口使用的DLL实现方法

这里我们将mobilenet作为一个私有类封装在一个dll里，.h文件声明方式如下：

// 下列 ifdef 块是创建使从 DLL 导出更简单的
// 宏的标准方法。此 DLL 中的所有文件都是用命令行上定义的DL_MOBILENETV3_EXPORTS
// 符号编译的。在使用此 DLL 的
// 任何其他项目上不应定义此符号。这样，源文件中包含此文件的任何其他项目都会将
// DL_MOBILENETV3_API 函数视为是从 DLL 导入的，而此 DLL 则将用此宏定义的
// 符号视为是被导出的。
#ifdef DL_MOBILENETV3_EXPORTS
#define DL_MOBILENETV3_API __declspec(dllexport)
#else
#define DL_MOBILENETV3_API __declspec(dllimport)
#endif

#pragma once
#include<iostream>
#include<opencv2/opencv.hpp>

namespace nsImageRecongition
{
	class DL_MOBILENETV3_API imageRecongition
	{
	public:
		imageRecongition();
		~imageRecongition();

		bool initial(const std::string& filePath);
		int  brokenRecongition(const cv::Mat& inputImg);

	private:

		class mobilenetv3;
		mobilenetv3 *mobilenet;
	};
}

-.cpp的实现方法如下：

#include "stdafx.h"
#include "DL_mobilenetv3.h"
#include "mobilenetv3.h"

nsImageRecongition::imageRecongition::imageRecongition()
{
	mobilenet = new mobilenetv3();
}

nsImageRecongition::imageRecongition::~imageRecongition()
{
	delete mobilenet;
}

bool nsImageRecongition::imageRecongition::initial(const std::string& filePath)
{
	return mobilenet->initial(filePath);
}

int nsImageRecongition::imageRecongition::brokenRecongition(const cv::Mat& inputImg)
{
	return mobilenet->predict(inputImg);
}

• mobilenet.h的声明方式如下:

#pragma once
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <cmath>

#include <opencv2/opencv.hpp>
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/imgproc/imgproc_c.h"
#include "opencv2/core/internal.hpp"

#include "DL_mobilenetv3.h"


#define CHECK(status) \
    do\
    {\
        auto ret = (status);\
        if (ret != 0)\
        {\
            std::cerr << "Cuda failure: " << ret << std::endl;\
            abort();\
        }\
    } while (0)


using namespace nvinfer1;

//using namespace nsImageRecongition;


class nsImageRecongition::imageRecongition::mobilenetv3
{

public:

	mobilenetv3();
	~mobilenetv3();

	bool initial(const std::string& filePath);

	int predict(const cv::Mat& inputImg);

private:

	int INPUT_H;

	int INPUT_W;
	int OUTPUT_SIZE;
	int BS;

	const char* INPUT_BLOB_NAME;
	const char* OUTPUT_BLOB_NAME;

	float* data;
	float* prob;

	IExecutionContext* context;
	ICudaEngine* engine;
	IRuntime* runtime;

	Logger gLogger;

	char *trtModelStream;
	size_t size;

	void doInference(IExecutionContext& context, float* input, float* output, int batchSize);
};

4. 测试

编译工程会生成相应的dll和exe文件，然后直接运行exe文件即可调用mobilenet模型进行推理预测。

双击Batch_Test.exe文件，然后弹出如下控制台窗口即说明封装调用成功