1 onnx模型的保存
在网络训练结束之后，通常会将模型的权重参数保存到.pth或.pt文件中，如果部署环境中有pytorch，那么直接新建一个模型类对象，然后导入权重参数即可，但如果部署环境中只有OpenCV，没有pytorch，那么该如何部署呢？
答：先在训练环境中将.pth文件转成onnx文件，再将onnx文件部署到最终的环境中，导出onnx文件的命令为torch.onnx.export即可。
import torch
import torchvision as tv
from torch.utils.data import DataLoader
class CNN_Mnist(torch.nn.Module):
def __init__(self):
super(CNN_Mnist, self).__init__()
self.cnn_layers = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1, stride=1),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.ReLU(),
torch.nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, padding=1, stride=1),
torch.nn.MaxPool2d(kernel_size=2, stride=2),
torch.nn.ReLU()
)
self.fc_layers = torch.nn.Sequential(
torch.nn.Linear(7*7*32, 200),
torch.nn.ReLU(),
torch.nn.Linear(200, 100),
torch.nn.ReLU(),
torch.nn.Linear(100, 10),
torch.nn.LogSoftmax(dim=1)
)
def forward(self, x):
out = self.cnn_layers(x)
out = out.view(-1, 7*7*32)
out = self.fc_layers(out)
return out
def train_and_test():
model = CNN_Mnist().cuda()
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for s in range(5):
print("run in epoch : %d" % s)
for i, (x_train, y_train) in enumerate(train_dl):
x_train = x_train.cuda()
y_train = y_train.cuda()
y_pred = model.forward(x_train)
train_loss = loss(y_pred, y_train)
if (i + 1) % 100 == 0:
print(i + 1, train_loss.item())
optimizer.zero_grad()
train_loss.backward()
optimizer.step()
torch.save(model.state_dict(), './cnn_mnist_model.pt')
model.eval()
total = 0
correct_count = 0
for test_images, test_labels in test_dl:
pred_labels = model(test_images.cuda())
predicted = torch.max(pred_labels, 1)[1]
correct_count += (predicted == test_labels.cuda()).sum()
total += len(test_labels)
print("total acc : %.2f\n"%(correct_count / total))
if __name__ == '__main__':
# 数据预处理方法
transform = tv.transforms.Compose([tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5,), (0.5,)),
])
# 数据集
train_ts = tv.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_ts = tv.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# 数据集导入器
train_dl = DataLoader(train_ts, batch_size=32, shuffle=True, drop_last=False)
test_dl = DataLoader(test_ts, batch_size=64, shuffle=True, drop_last=False)
# 训练与测试
train_and_test()
# 模型的导入
model = CNN_Mnist()
model.load_state_dict(torch.load('cnn_mnist_model.pt'))
"""下面演示将模型转化为onnx格式文件"""
# 要把模型切换到评估状态，这样可以让某些层（如drop_out失效）
model.eval()
# 随机一个输入张量，这个张量的作用是告诉onnx框架，输入张量的shape是什么样的
dummy_input = torch.randn(1, 1, 224, 224)
# 导出onnx文件
torch.onnx.export(model, (dummy_input), 'cnn_mnistorch.onnx', verbose=True) # onnx文件可以通过netron查看结构
# verbose=True，则打印一些转换日志，并且onnx文件中会包含doc_string，即用于说明模型的文档字符串
这段程序执行完成后，就会在当前目录下得到一个名为“cnn_mnistorch.onnx”的文件
2 在OpenCV中调用onnx模型文件
OpenCV调用onnx文件也很方便，只需要cv.dnn.readNetFromONNX即可，剩下的部分就是调用OpenCV了
import cv2 as cv
import numpy as np
def mnist_onnx_demo():
# 从onnx文件中读取网络
mnist_net = cv.dnn.readNetFromONNX("cnn_mnist.onnx")
# 读取图片并做相应的预处理
image = cv.imread("test.png")
gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
cv.imshow("input", gray)
blob = cv.dnn.blobFromImage(gray, 0.00392, (28, 28), (127.0)) / 0.5 # 这个API稍后会讲
"""上面对应图像在pytorch中的变换
tv.transforms.Compose([tv.transforms.ToTensor(),
tv.transforms.Normalize((0.5,), (0.5,)),])"""
print(blob.shape)
# 前向传播，OpenCV的dnn模块，前向传播需要先设置网络的输入
mnist_net.setInput(blob)
result = mnist_net.forward()
# 后处理
pred_label = np.argmax(result, 1)
print("predit label : %d" % pred_label)
cv.waitKey(0)
cv.destroyAllWindows()
# 整个过程没有调用pytorch框架，也没使用模型对应的类，也就是说，onnx文件可以摆脱对框架的依赖
if __name__ == '__main__':
mnist_onnx_demo()
输出：
(1, 1, 28, 28)
predit label : 3
可以看到，部署的过程中，完全拜托了对pytorch框架的依赖。
这里面比较重要的是cv2.dnn.blobFromImage这个函数，它的作用是将图像转化为网络模型的输入，API如下：
cv2.dnn.blobFromImage(image[, scalefactor[, size[, mean[, swapRB[, crop[, ddepth]]]]]])
作用：
对图像进行预处理，包括减均值，比例缩放，裁剪，交换通道等，返回一个4通道的blob(blob可以简单理解为一个N维的数组，用于神经网络的输入)
参数：
image:输入图像（1、3或者4通道）
可选参数
scalefactor:图像各通道数值的缩放比例
size:图像要转化成的空间尺寸,如size=(200,300)表示高h=300,宽w=200，相当于resize
mean:用于各通道减去的值，以降低光照的影响(e.g. image为BGR3通道的图像，mean=[104.0, 177.0, 123.0],表示b通道的值-104，g-177,r-123)
swapRB:交换RB通道，默认为False.(cv2.imread读取的是彩图是bgr通道)
crop:图像裁剪,默认为False.
当值为True时，先按保持原来的高宽比缩放，直到其中一条边等于对应方向的长度，另一条边大于对应方向长度，然后从中心裁剪成size尺寸；
如果值为False，则不管高宽比，直接缩放成指定尺寸（即size参数）。
e.g.原图(300, 200)，目标尺寸(400, 300)
若crop=True， (300, 200) --resize-->(450, 300)--crop-->(400, 300)
若crop=False， (300, 200) --resize-->(400, 300)
ddepth:输出的图像深度，可选CV_32F 或者 CV_8U.
