@[TOC]

摘要

本例提取了植物幼苗数据集中的部分数据做数据集，数据集共有12种类别，演示如何使用pytorch版本的mobilenetv2图像分类模型实现分类任务。将训练的模型转为onnx，实现onnx的推理，然后再将onnx转为TensorRT，并实现推理。

通过本文你和学到：

​ 1、如何从torchvision.models调用mobilenetv2模型？

​ 2、如何自定义数据集加载方式？

​ 3、如何使用Cutout数据增强？

​ 4、如何使用Mixup数据增强。

​ 5、如何实现训练和验证。

​ 6、如何使用余弦退火调整学习率。

​ 7、如何载入训练的模型进行预测。

​ 8、pytorch转onnx，并实现onnx推理。

​ 9、onnx转TensorRT，并实现TensorRT的推理。

希望通过这篇文章，能让大家对图像的分类和模型的部署有个清晰的认识。

mobilenetv2简介

mobileNetV2是对mobileNetV1的改进，是一种轻量级的神经网络。mobileNetV2保留了V1版本的深度可分离卷积，增加了线性瓶颈（Linear Bottleneck）和倒残差（Inverted Residual）。

MobileNetV2的模型如下图所示，其中t为瓶颈层内部升维的倍数，c为特征的维数，n为该瓶颈层重复的次数，s为瓶颈层第一个conv的步幅。

除第一层外，整个网络中使用恒定的扩展率。 在实验中，发现在 5 到 10 之间的扩展率会导致几乎相同的性能曲线，较小的网络在扩展率稍低的情况下效果更好，而较大的网络在扩展率较大的情况下性能稍好。

MobileNetV2主要使用 6 的扩展因子应用于输入张量的大小。 例如，对于采用 64 通道输入张量并产生具有 128 通道的张量的瓶颈层，则中间扩展层为 64×6 = 384 通道。

线性瓶颈

对于mobileNetV1的深度可分离卷积而言， 宽度乘数压缩后的M维空间后会通过一个非线性变换ReLU，根据ReLU的性质，输入特征若为负数，该通道的特征会被清零，本来特征已经经过压缩，这会进一步损失特征信息；若输入特征是正数，经过激活层输出特征是还原始的输入值，则相当于线性变换。

瓶颈层的具体结构如下表所示。输入通过1的conv+ReLU层将维度从k维增加到tk维，之后通过3×3conv+ReLU可分离卷积对图像进行降采样（stride>1时），此时特征维度已经为tk维度，最后通过1*1conv（无ReLU）进行降维，维度从tk降低到k维。

倒残差

残差块已经在ResNet中得到证明，有助于提高精度构建更深的网络，所以mobileNetV2也引入了类似的块。经典的残差块（residual block）的过程是：1x1(降维)–>3x3(卷积)–>1x1(升维)， 但深度卷积层（Depthwise convolution layer）提取特征限制于输入特征维度，若采用残差块，先经过1x1的逐点卷积（Pointwise convolution）操作先将输入特征图压缩，再经过深度卷积后，提取的特征会更少。所以mobileNetV2是先经过1x1的逐点卷积操作将特征图的通道进行扩张，丰富特征数量，进而提高精度。这一过程刚好和残差块的顺序颠倒，这也就是倒残差的由来：1x1(升维)–>3x3(dw conv+relu)–>1x1(降维+线性变换)。

结合上面对线性瓶颈和倒残差的理解，我绘制了Block的结构图。如下图：

ONNX

ONNX，全称：Open Neural Network Exchange（ONNX，开放神经网络交换），是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。

ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM 等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX。—维基百科

onnx模型可以看作是模型转化之间的中间模型，同时也是支持做推理的。一般来说，onnx的推理速度要比pytorch快上一倍。

TensorRT

TensorRT是英伟达推出的一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。

TensorRT 是一个C++库，从 TensorRT 3 开始提供C++ API和Python API，主要用来针对 NVIDIA GPU进行 高性能推理（Inference）加速。

项目结构

MobileNetV2_demo
├─data
│  ├─test
│  └─train
│      ├─Black-grass
│      ├─Charlock
│      ├─Cleavers
│      ├─Common Chickweed
│      ├─Common wheat
│      ├─Fat Hen
│      ├─Loose Silky-bent
│      ├─Maize
│      ├─Scentless Mayweed
│      ├─Shepherds Purse
│      ├─Small-flowered Cranesbill
│      └─Sugar beet
├─dataset
│  └─dataset.py
├─train.py
├─test_torch.py
├─torch2onnx.py
├─test_onnx.py
├─onnx2trt.py
└─test_trt.py

训练

数据增强Cutout和Mixup

为了提高成绩我在代码中加入Cutout和Mixup这两种增强方式。实现这两种增强需要安装torchtoolbox。安装命令：

pip install torchtoolbox

Cutout实现，在transforms中。

from torchtoolbox.transform import Cutout

# 数据预处理

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    Cutout(),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

])

Mixup实现，在train方法中。需要导入包：from torchtoolbox.tools import mixup_data, mixup_criterion

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
        data, labels_a, labels_b, lam = mixup_data(data, target, alpha)
        optimizer.zero_grad()
        output = model(data)
        loss = mixup_criterion(criterion, output, labels_a, labels_b, lam)
        loss.backward()
        optimizer.step()
        print_loss = loss.data.item()

导入包

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import SeedlingData
from torch.autograd import Variable
from torchvision.models import mobilenet_v2
from torchtoolbox.tools import mixup_data, mixup_criterion
from torchtoolbox.transform import Cutout

设置全局参数

设置学习率、BatchSize、epoch等参数，判断环境中是否存在GPU，如果没有则使用CPU。建议使用GPU，CPU太慢了。由于mobilenetv2模型很小，4G显存的GPU就可以设置BatchSize为16。

# 设置全局参数
modellr = 1e-4
BATCH_SIZE = 16
EPOCHS = 300
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

图像预处理与增强

数据处理比较简单，加入了Cutout、做了Resize和归一化。对于Normalize和std的值，这个一般是通用的设置，而且在后面的测试中要保持一致。

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    Cutout(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

读取数据

将数据集解压后放到data文件夹下面，如图：

然后我们在dataset文件夹下面新建 init.py和dataset.py，在datasets.py文件夹写入下面的代码：

# coding:utf8
import os
from PIL import Image
from torch.utils import data
from torchvision import transforms as T
from sklearn.model_selection import train_test_split
 
Labels = {'Black-grass': 0, 'Charlock': 1, 'Cleavers': 2, 'Common Chickweed': 3,
          'Common wheat': 4, 'Fat Hen': 5, 'Loose Silky-bent': 6, 'Maize': 7, 'Scentless Mayweed': 8,
          'Shepherds Purse': 9, 'Small-flowered Cranesbill': 10, 'Sugar beet': 11}
 
 
class SeedlingData (data.Dataset):
 
    def __init__(self, root, transforms=None, train=True, test=False):
        """
        主要目标： 获取所有图片的地址，并根据训练，验证，测试划分数据
        """
        self.test = test
        self.transforms = transforms
 
        if self.test:
            imgs = [os.path.join(root, img) for img in os.listdir(root)]
            self.imgs = imgs
        else:
            imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]
            imgs = []
            for imglable in imgs_labels:
                for imgname in os.listdir(imglable):
                    imgpath = os.path.join(imglable, imgname)
                    imgs.append(imgpath)
            trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)
            if train:
                self.imgs = trainval_files
            else:
                self.imgs = val_files
 
    def __getitem__(self, index):
        """
        一次返回一张图片的数据
        """
        img_path = self.imgs[index]
        img_path=img_path.replace("\\",'/')
        if self.test:
            label = -1
        else:
            labelname = img_path.split('/')[-2]
            label = Labels[labelname]
        data = Image.open(img_path).convert('RGB')
        data = self.transforms(data)
        return data, label
 
    def __len__(self):
        return len(self.imgs)

说一下代码的核心逻辑：

第一步 建立字典，定义类别对应的ID,用数字代替类别。

第二步 在__init__里面编写获取图片路径的方法。测试集只有一层路径直接读取，训练集在train文件夹下面是类别文件夹，先获取到类别，再获取到具体的图片路径。然后使用sklearn中切分数据集的方法，按照7：3的比例切分训练集和验证集。

第三步 在__getitem__方法中定义读取单个图片和类别的方法，由于图像中有位深度32位的，所以我在读取图像的时候做了转换。

然后我们在train.py调用SeedlingData读取数据 ，记着导入刚才写的dataset.py(from dataset.dataset import SeedlingData)

dataset_train = SeedlingData('data/train', transforms=transform, train=True)
dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)
# 读取数据
print(dataset_train.imgs)

# 导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

设置模型

• 设置loss函数为nn.CrossEntropyLoss()。
• 设置模型为mobilenet_v3_large，预训练设置为true，num_classes设置为12。
• 优化器设置为adam。
• 学习率调整策略选择为余弦退火。

# 实例化模型并且移动到GPU
criterion = nn.CrossEntropyLoss()
#criterion = SoftTargetCrossEntropy()
model_ft = mobilenet_v2(pretrained=True)
print(model_ft)
num_ftrs = model_ft.classifier[1].in_features
model_ft.classifier[1] = nn.Linear(num_ftrs, 12,bias=True)
model_ft.to(DEVICE)
print(model_ft)
# 选择简单暴力的Adam优化器，学习率调低
optimizer = optim.Adam(model_ft.parameters(), lr=modellr)
cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer,T_max=20,eta_min=1e-9)
#model_ft=torch.load("model_33_0.952.pth")

定义训练和验证函数

# 定义训练过程
alpha=0.2
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    sum_loss = 0
    total_num = len(train_loader.dataset)
    print(total_num, len(train_loader))
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
        data, labels_a, labels_b, lam = mixup_data(data, target, alpha)
        optimizer.zero_grad()
        output = model(data)
        loss = mixup_criterion(criterion, output, labels_a, labels_b, lam)
        loss.backward()
        optimizer.step()
        lr = optimizer.state_dict()['param_groups'][0]['lr']
        print_loss = loss.data.item()
        sum_loss += print_loss
        if (batch_idx + 1) % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR:{:.9f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item(),lr))
    ave_loss = sum_loss / len(train_loader)
    print('epoch:{},loss:{}'.format(epoch, ave_loss))

ACC=0
# 验证过程
def val(model, device, test_loader):
    global ACC
    model.eval()
    test_loss = 0
    correct = 0
    total_num = len(test_loader.dataset)
    print(total_num, len(test_loader))
    with torch.no_grad():
        for data, target in test_loader:
            data, target = Variable(data).to(device), Variable(target).to(device)
            output = model(data)
            loss = criterion(output, target)
            _, pred = torch.max(output.data, 1)
            correct += torch.sum(pred == target)
            print_loss = loss.data.item()
            test_loss += print_loss
        correct = correct.data.item()
        acc = correct / total_num
        avgloss = test_loss / len(test_loader)
        print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            avgloss, correct, len(test_loader.dataset), 100 * acc))
        if acc > ACC:
            torch.save(model_ft, 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')
            ACC = acc


# 训练

for epoch in range(1, EPOCHS + 1):
    train(model_ft, DEVICE, train_loader, optimizer, epoch)
    cosine_schedule.step()
    val(model_ft, DEVICE, test_loader)

完成上面的代码就可以运行了！待训练完成后，我们就可以进行下一步测试的工作了。

测试

这里介绍一种通用的测试方法，通过自己手动加载数据集然后做预测，具体操作如下：

测试集存放的目录如下图：

第一步 定义类别，这个类别的顺序和训练时的类别顺序对应，一定不要改变顺序！！！！

第二步 定义transforms，transforms和验证集的transforms一样即可，别做数据增强。

第三步 加载model，并将模型放在DEVICE里，

第四步 读取图片并预测图片的类别，在这里注意，读取图片用PIL库的Image。不要用cv2，transforms不支持。

import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os
classes = ('Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed',
           'Common wheat','Fat Hen', 'Loose Silky-bent',
           'Maize','Scentless Mayweed','Shepherds Purse','Small-flowered Cranesbill','Sugar beet')
transform_test = transforms.Compose([
         transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
 
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
 
path='data/test/'
testList=os.listdir(path)
for file in testList:
        img=Image.open(path+file)
        img=transform_test(img)
        img.unsqueeze_(0)
        img = Variable(img).to(DEVICE)
        out=model(img)
        # Predict
        _, pred = torch.max(out.data, 1)
        print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))

运行结果：

模型转化及推理

转onnx

模型可视化使用了netron，安装命令：

pip install netron

新建torch2onnx.py脚本，插入代码：

import torch
import torchvision
from torch.autograd import Variable
import netron
print(torch.__version__)
# torch  -->  onnx
input_name = ['input']
output_name = ['output']
input = Variable(torch.randn(1, 3, 224, 224)).cuda()
model = torch.load('model_59_0.938.pth', map_location="cuda:0")
torch.onnx.export(model, input, 'model_onnx.onnx',opset_version=12, input_names=input_name, output_names=output_name, verbose=True)
# 模型可视化
netron.start('model_onnx.onnx')

导入需要的包。

打印pytorch版本。

定义input_name和output_name变量。

定义输入格式。

加载pytorch模型。

导出onnx模型，这里注意一下参数opset_version在8.X版本中设置为13，在7.X版本中设置为12。

onnx推理

onnx推理需要安装onnx包文件，命令如下：

pip install onnx
pip install onnxruntime-gpu

新建test_onnx.py，插入下面的代码：

import os, sys
import time
sys.path.append(os.getcwd())
import onnxruntime
import numpy as np
import torchvision.transforms as transforms
from PIL import Image

导入包

def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
image = Image.open('11.jpg') # 289
img = get_test_transform()(image)
img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(img.mean(), img.std()))
img =  np.array(img)

定义get_test_transform函数，实现图像的归一化和resize。

读取图像。

对图像做resize和归一化。

增加一维batchsize。

将图片转为数组。

onnx_model_path = "model_onnx.onnx"
##onnx测试
session = onnxruntime.InferenceSession(onnx_model_path,providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'])
#compute ONNX Runtime output prediction
inputs = {session.get_inputs()[0].name: img}
time3=time.time()
outs = session.run(None, inputs)[0]
y_pred_binary = np.argmax(outs, axis=1)
print("onnx prediction", y_pred_binary[0])
time4=time.time()
print(time4-time3)

定义onnx_model_path模型的路径。

加载onnx模型。

定义输入。

执行推理。

获取预测结果。

到这里onnx推理就完成了，是不是很简单！！！

转TensorRT

新建onnx2trt.py,插入代码

import tensorrt as trt


def build_engine(onnx_file_path,engine_file_path,half=False):
    """Takes an ONNX file and creates a TensorRT engine to run inference with"""
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    config.max_workspace_size = 4 * 1 << 30#这决定了可用的内存量 
    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, logger)
    if not parser.parse_from_file(str(onnx_file_path)):
        raise RuntimeError(f'failed to load ONNX file: {onnx_file_path}')
    half &= builder.platform_has_fast_fp16
    if half:
        config.set_flag(trt.BuilderFlag.FP16)
    with builder.build_engine(network, config) as engine, open(engine_file_path, 'wb') as t:
        t.write(engine.serialize())
    return engine_file_path
if __name__ =="__main__":
    onnx_path1 = 'model_onnx.onnx'
    engine_path = 'model_trt.engine'
    build_engine(onnx_path1,engine_path,True)

build_engine函数共有三个参数：

onnx_file_path：onnx模型的路径。

engine_file_path：TensorRT模型的路径。

half：是否使用单精度。

单精度的模型速度更快，所以我选择使用单精度。

build_engine函数中：

​ 实现一个日志接口，TensorRT通过该接口报告错误、警告和信息消息。

​ 创建TensorRT的builder生成器。

​ 创建builder的Config。Config有许多属性，您可以设置这些属性来控制网络运行的精度，以及自校正参数，一个特别重要的属性是最大工作空间大小。

​ 创建network。

​ 创建parser。

​ 如果存在onnx模型，则载入。

​ 如果half为true，则在builder和config中分别设置platform_has_fast_fp16和trt.BuilderFlag.FP16。

​ 接下来，将engine序列化后保存。推理的时候在反序列化载入。

TensorRT推理

新建test_trt,py文件，插入代码：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import torchvision.transforms as transforms
from PIL import Image

导入需要的包。

def load_engine(engine_path):
    # TRT_LOGGER = trt.Logger(trt.Logger.WARNING)  # INFO
    TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
    with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

# 2. 读取数据，数据处理为可以和网络结构输入对应起来的的shape，数据可增加预处理
def get_test_transform():
    return transforms.Compose([
        transforms.Resize([224, 224]),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

定义load_engine函数和get_test_transform函数。

load_engine用于加载TensorRT模型。

get_test_transform实现图像的resize和归一化。

image = Image.open('data/test/00c47e980.png') # 289
image = get_test_transform()(image)
image = image.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(image.mean(), image.std()))
image =  np.array(image)

图片的预处理，和onnx一样，最后转为数组。

path = 'model_trt.engine'
# 1. 建立模型，构建上下文管理器
engine = load_engine(path)
context = engine.create_execution_context()

# 3.分配内存空间，并进行数据cpu到gpu的拷贝
# 动态尺寸，每次都要set一下模型输入的shape，0代表的就是输入，输出根据具体的网络结构而定，可以是0,1,2,3...其中的某个头。
context.set_binding_shape(0, image.shape)
d_input = cuda.mem_alloc(image.nbytes)  # 分配输入的内存。
output_shape = context.get_binding_shape(1)
buffer = np.empty(output_shape, dtype=np.float32)
d_output = cuda.mem_alloc(buffer.nbytes)  # 分配输出内存。
cuda.memcpy_htod(d_input, image)
bindings = [d_input, d_output]

# 4.进行推理，并将结果从gpu拷贝到cpu。
context.execute_v2(bindings)  # 可异步和同步
cuda.memcpy_dtoh(buffer, d_output)
output = buffer.reshape(output_shape)
y_pred_binary = np.argmax(output, axis=1)
print(y_pred_binary[0])

输出结果：

完整代码下载：
https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/81304254