TVM-LeNet
本文完成了自己搭建、训练和打包LeNet模型,并用TVM对其进行简单的优化,比较两者的异同,发现TVM的计算图更为简捷,准确率二者几乎相同,并且TVM的推理速度是普通模型的2倍左右。
版本:TVM 0.7.dev1 ; pytorch 1.6.0 ; GTX1650 , cuda10.2
1.LeNet模型搭建:
LeNet-5是1998年Yann Lecun教授发表的一种用于灰度识别的网络,大致的网络结构如下:
如图,LeNet输入为32*32的灰度图片,2个卷积层,两个最大池化层,3个全连接层,我是使用pytorch进行网络的搭建,模型部分代码如下:
## lenet_model.ipynb
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 5)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, 5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32*5*5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28)
x = self.pool1(x) # output(16, 14, 14)
x = F.relu(self.conv2(x)) # output(32, 10, 10)
x = self.pool2(x) # output(32, 5, 5)
x = x.view(-1, 32*5*5) # output(32*5*5)
x = F.relu(self.fc1(x)) # output(120)
x = F.relu(self.fc2(x)) # output(84)
x = self.fc3(x) # output(10)
return x
pytorch框架为模型搭建提供了很方便的方法。比如第一个卷积层self.conv1 = nn.Conv2d(3, 16, 5)
,其函数定义如下:
由于我这次是利用CIFAR-10数据集进行训练和推理,该数据集是RGB3色的,因此卷积层的in_channels设置为3,out_channels的设置并不严格,我们设置的比论文中小一些,设为16,卷积核大小设置为5x5,步长默认为1,padding默认补零,也就是说,经过第一个卷积层后,会提取到包含不同特征的16张feature_map,并且图片大小变为28x28。即output为(16,28,28)
下面的第一个池化层nn.MaxPool2d(2,2)
的定义如下:
2x2的池化层会将输入特征图的尺寸减半,因此输出为(16,14,14),以此类推可以分析出下面一组卷积池化层。
卷积池化结束后开始进行全连接,经最后一次池化,得到输出的shape为(32,5,5),全连接层的输入必须是一维的tensor,因此利用函数x.view(-1, 32*5*5)
将其展平,该函数定义为:
可按此方法分析剩下两个全连接层,并最终得到一个(1,10)的输出,张量中每个元素的index与训练集中种类的index相同,元素的值代表着图片是这种类的可能性。
这样一个LeNet-5模型就搭建好了。
2.模型训练:
##lenet_train.ipynb
import Ipynb_importer
import torch
import torchvision
import torch.nn as nn
from lenet_model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
transform函数用于对之后的图片进行预处理。
从CIFAR10下载训练集和数据集:
# 10000张验证图片
# 第一次使用时要将download设置为True才会自动去下载数据集
val_set = torchvision.datasets.CIFAR10(root='/home/ljs/Desktop/data', train=False,
download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=5000,
shuffle=False, num_workers=4)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
图片信息存储在val_image中,图片标签存储在val_label中。
由于训练时需要使用GPU,按照框架的规则,需要先对图片和标签数据进行处理,如下:
val_image=val_image.cuda()
val_label=val_label.cuda()
net = LeNet()
net=net.cuda()
loss_function = nn.CrossEntropyLoss() #损失函数为CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
数据准备好后就可以开始训练,训练部分的代码如下:
for epoch in range(5): # loop over the dataset multiple times
running_loss = 0.0
for step, data in enumerate(train_loader, start=0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
inputs=inputs.cuda()
labels=labels.cuda()
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if step % 500 == 499: # print every 500 mini-batches
with torch.no_grad():
outputs = net(val_image) # [batch, 10]
predict_y = torch.max(outputs, dim=1)[1]
accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0)
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, step + 1, running_loss / 500, accuracy))
running_loss = 0.0
训练五批,最终达到67%的准确率:
通过下面函数可以保存训练好的模型参数:
save_path = '/home/ljs/Desktop/Lenet.pth'
torch.save(net.state_dict(), save_path)
3.模型推理:
经过训练部分同样的图片预处理后,实例化模型,导入权重参数,即可进行前向推理:
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = LeNet()
net=net.cuda()
net.load_state_dict(torch.load('/home/ljs/Desktop/Lenet.pth'))
for im in val_image:
im=torch.unsqueeze(im, dim=0).cuda()
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1]
print(classes[int(predict)])
测试结果如下准确性还不错:
4.使用TVM进行推理:
## lenet_tvm.ipynb,需要导入的库
import Ipynb_importer
import torch
import torchvision.transforms as transforms
from PIL import Image
from lenet_model import LeNet
import torchvision
import matplotlib.pyplot as plt
import numpy as np
import tvm
import torchvision
from tvm import relay
import datetime
参考TVM代码阅读——pytorch这篇文章,利用TVM来优化编译第一节搭建的模型。遵循TVM中API的规则,需要对模型先做处理,利用scripted_model=torch.jit.trace(Net, input_data).eval()
这个函数,将模型转换为TopLevelTracedModule类,并且TVM接收的图像数据就是常规数据,不需要提前转换移动到GPU中,TVM会调用相应的编译器与GPU关联。
使用CIFAR10中的50000张照片进行测试:
val_set = torchvision.datasets.CIFAR10(root='/home/ljs/Desktop/data', train=True,
download=False, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=60000,
shuffle=False, num_workers=0)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
val_image.shape
input_data=torch.rand(1, 3, 32, 32) ##
int(val_label[1])
val_image.shape ## -->torch.Size([50000, 3, 32, 32])
导入模型并转换类型:
transform = transforms.Compose(
[transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')
Net=LeNet()
Net.load_state_dict(torch.load('/home/ljs/Desktop/Lenet.pth'))
scripted_model=torch.jit.trace(Net, input_data).eval()
type(scripted_model) ##-->torch.jit.TopLevelTracedModule
获取TVM需要的数据:
input_name = "input0"
shape_list = [(input_name, img_shape.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
选择后端和优化等级:
target = "cuda"
target_host = "llvm"
ctx = tvm.gpu(0)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, target_host=target_host, params=params)
构建relay graph,开始推理:
from tvm.contrib import graph_runtime
dtype = "float32"
m = graph_runtime.GraphModule(lib["default"](ctx))
for im in img:
im= np.expand_dims(im, 0)
# Set inputs
m.set_input(input_name, tvm.nd.array(im.astype(dtype)))
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)
top1_tvm = np.argmax(tvm_output.asnumpy()[0])
print(classes[top1_tvm])
截取推理得到的前前五张图片标签:
发现和原始模型推理的一致。
5.结果比较:
首先比较二者计算图,彩色图片为未经优化的模型的计算图,通过将pytorch的模型转换为onnx模式,用netron工具进行可视化;黑白图片为经优化后的计算图,通过
with open("/home/ljs/Desktop/opted-lenet-graph.txt","w") as f:
f.write(str(lib.graph_json))
将计算图导出,netron可以可视化json格式的计算图:
![在这里插入图片描述](https://i-blog.csdnimg.cn/blog_migrate/90808c659124a1d4f81ec264b5980649.png#pic_center)
从图中可以看出,由于TVM会进行算子融合等优化手段,将卷积层与池化层进行了融合,优化后的计算图的层数会更少。
再比较二者推理时间,两个模型对同样的50000张图片进行识别(这里其实不能这样利用训练集数据测试的,但cifar10只有10000张测试图片,为两种模型运行时间差距明显,使用50000张图片进行测试;10000张和50000笔者实验后,发现耗时比都是2.几,区别不大),并且计算准确率,结果如下:
模型 | 耗时(s) | 准确率 |
---|---|---|
LeNet | 18.381974 | 73.756% |
TVM-LeNet | 7.022352 | 73.754% |