摘要
本例提取了植物幼苗数据集中的部分数据做数据集，数据集共有 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。安装命令：
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 optimimport torchimport torch.nn as nnimport torch.nn.parallelimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsfrom dataset.dataset import SeedlingDatafrom torch.autograd import Variablefrom torchvision.models import mobilenet_v2from torchtoolbox.tools import mixup_data, mixup_criterionfrom torchtoolbox.transform import Cutout
复制代码
设置全局参数
设置学习率、BatchSize、epoch 等参数，判断环境中是否存在 GPU，如果没有则使用 CPU。建议使用 GPU，CPU 太慢了。由于 mobilenetv2 模型很小，4G 显存的 GPU 就可以设置 BatchSize 为 16。
# 设置全局参数modellr = 1e-4BATCH_SIZE = 16EPOCHS = 300DEVICE = 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:utf8import osfrom PIL import Imagefrom torch.utils import datafrom torchvision import transforms as Tfrom 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)
复制代码
设置模型
# 实例化模型并且移动到GPUcriterion = nn.CrossEntropyLoss()#criterion = SoftTargetCrossEntropy()model_ft = mobilenet_v2(pretrained=True)print(model_ft)num_ftrs = model_ft.classifier[1].in_featuresmodel_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.2def 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.distributedimport torchvision.transforms as transformsfrom PIL import Imagefrom torch.autograd import Variableimport osclasses = ('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，安装命令：
新建 torch2onnx.py 脚本，插入代码：
import torchimport torchvisionfrom torch.autograd import Variableimport netronprint(torch.__version__)# torch --> onnxinput_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 onnxpip install onnxruntime-gpu
复制代码
新建 test_onnx.py，插入下面的代码：
import os, sysimport timesys.path.append(os.getcwd())import onnxruntimeimport numpy as npimport torchvision.transforms as transformsfrom 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') # 289img = get_test_transform()(image)img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224print("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 predictioninputs = {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_pathif __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 trtimport pycuda.driver as cudaimport pycuda.autoinitimport numpy as npimport torchvision.transforms as transformsfrom 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') # 289image = get_test_transform()(image)image = image.unsqueeze_(0) # -> NCHW, 1,3,224,224print("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
评论