LibTorch实战一：MNIST的python代码
目录
一、下载MNIST-demo的cpp、python版本代码
二、五分钟读懂pytorch代码
三、下载MNIST数据集、训练模型
四、模型序列化、可视化分析
本文借用mnist这个相对简易深度学习任务来开始讲解libtorch如何部署模型。因此，这是一个如何编写libtorch代码的实战教程。
一、下载MNIST-demo的cpp、python版本代码
进入链接：https://github.com/pytorch 如下图：打开example
将上述example下的code下到本地，在根目录可以看到mnist的python版本：main.py；打开cpp/mnist，可以看到mnist的c++版本(libtorch实现的)：mnist.cpp。以上两个代码文件下文会用到。
mnist是一个手写数字识别的demo，如下图，快速了解下原理。
二、五分钟读懂pytorch代码
读透上述python版代码。打开main.py，我们知道，输入图片大小为28*28。下面是整个MNISTpython官方代码的解读，是训练代码，需要下载数据集，已经很详细注释。如下图，执行之
后会自动下载手写数字数据，然后开始训练。(torch1.13.1)
注：如果你不能自动下载数据集，请看第三章。
main.py:
from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
class Net(nn.Module):
def __init__(self): # self指的是类实例对象本身(注意：不是类本身)。
# self不是关键词
# super 用于继承，https://www.runoob.com/python/python-func-super.html
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
# input:28*28
x = self.conv1(x) # -> (28 - 3 + 1 = 26),26*26*32
x = F.relu(x)
# input:26*26*32
x = self.conv2(x) # -> (26 - 3 + 1 = 24),24*24*64
# input:24*24*64
x = F.relu(x)
x = F.max_pool2d(x, 2)# -> 12*12*64 = 9216
x = self.dropout1(x) #不改变维度
x = torch.flatten(x, 1) # 9216*1
# w = 128*9216
x = self.fc1(x) # -> 128*1
x = F.relu(x)
x = self.dropout2(x)
# w = 10*128
x = self.fc2(x) # -> 10*1
output = F.log_softmax(x, dim=1) # softmax归一化
return output
def train(args, model, device, train_loader, optimizer, epoch):
# 在使用pytorch构建神经网络的时候，训练过程中会在程序上方添加一句model.train()，
# 作用是启用batch normalization和drop out。
# 测试过程中会使用model.eval()，这时神经网络会沿用batch normalization的值，并不使用drop out。
model.train()
# 可以查看下卷积核的参数尺寸
#model.conv1.weight.shape torch.Size([32, 1, 3, 3] 即：32个2D卷积核
#model.conv2.weight.shape torch.Size([64, 32, 3, 3]) 即：64个3D卷积核
for batch_idx, (data, target) in enumerate(train_loader):
# train_loader.dataset.data.shape
# Out[9]: torch.Size([60000, 28, 28])
# batch_size:64
# data:64个样本输入,torch.Size([64, 1, 28, 28])
# target: 64个label,torch.Size([64])
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
# output:torch.Size([64, 10])
output = model(data)
# 类似于交叉熵
# reference: https://blog.csdn.net/qq_22210253/article/details/85229988
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
# 我们打印一个卷积核参数看看
# print(model.conv2._parameters)
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
if args.dry_run:
break
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
parser.add_argument('--batch-size', type=int, default=16, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=2, metavar='N',
help='number of epochs to train (default: 14)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--dry-run', action='store_true', default=False,
help='quickly check a single pass')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=True,
help='For Saving the current Model')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
train_kwargs = {'batch_size': args.batch_size}
test_kwargs = {'batch_size': args.test_batch_size}
if use_cuda:
cuda_kwargs = {'num_workers': 1,
'pin_memory': True, # 锁页内存，可以加快内存到显存的速度
'shuffle': True}
train_kwargs.update(cuda_kwargs)
test_kwargs.update(cuda_kwargs)
# torchvision.transforms是pytorch中的图像预处理包。一般用Compose把多个步骤整合到一起
#
transform = transforms.Compose([
transforms.ToTensor(), # (H x W x C)、[0, 255] -> (C x H x W)、[0.0, 1.0]
transforms.Normalize((0.1307,), (0.3081,)) # 数据的归一化
])
dataset1 = datasets.MNIST('data', train=True, download=False,
transform=transform)
dataset2 = datasets.MNIST('data', train=False,
transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
# 固定步长衰减
# reference: https://zhuanlan.zhihu.com/p/93624972
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(model, device, test_loader)
scheduler.step()
if args.save_model:
#torch.save(model.state_dict(), "pytorch_mnist.pt")
torch.save(model, "pytorch_mnist.pth")
if __name__ == '__main__':
main()三、下载MNIST数据集、训练模型
如何下载数据集：网上一堆教程，似乎没有什么用，讲道理是pytorch底层自动取固定网站下载，并在本地新建文件夹，存在如下如data相对路径，
我翻看底层，下个数据，xxx还要校验MD5码，简直有X。也就是说你通过下载别人分享的压缩包，pytorch是不能加载的，我x！不过后来我找到了.pt格式的数据集，原来pytorch读取压缩包.gz格式的数据后，再将其转为.pt格式的文件，存在processed文件夹中，如下图：
下载连接：https://github.com/MorvanZhou/PyTorch-Tutorial/tree/master/tutorial-contents-notebooks/mnist/processed
下载之后放到上述processed文件夹中就行，原始压缩包格式的数据集就不用管了，.pt格式的数据是可以拷贝共享给别人的。最终工程目录如下图：
运行main.py，执行python代码训练、保存模型文件pytorch_mnist.pth，测试打印信息如下：
......
184 Train Epoch: 2 [54400/60000 (91%)] Loss: 0.003272
185 Train Epoch: 2 [55040/60000 (92%)] Loss: 0.236524
186 Train Epoch: 2 [55680/60000 (93%)] Loss: 0.087931
187 Train Epoch: 2 [56320/60000 (94%)] Loss: 0.013646
188 Train Epoch: 2 [56960/60000 (95%)] Loss: 0.027721
189 Train Epoch: 2 [57600/60000 (96%)] Loss: 0.100714
190 Train Epoch: 2 [58240/60000 (97%)] Loss: 0.155445
191 Train Epoch: 2 [58880/60000 (98%)] Loss: 0.113110
192 Train Epoch: 2 [59520/60000 (99%)] Loss: 0.039872
193
194 Test set: Average loss: 0.0398, Accuracy: 9864/10000 (99%)这里再给一个在pytorch、python-opencv环境中的模型测试代码。代码中读取上述训练保存的模型文件，然后推理图像，图像下面随便给了两个小图。
示例图：
infer.py
from main import Net
from torchvision import datasets, transforms
from PIL import Image
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('pytorch_mnist.pth') # 加载模型
model = model.to(device)
model.eval() # 把模型转为test模式
img = cv2.imread("9.jpg", 0) # 读取要预测的灰度图片
cv2.imshow("img", img)
cv2.waitKey(100)
img = Image.fromarray(img)
trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img = trans(img)
img = img.unsqueeze(0) # 图片扩展多一维,[batch_size,通道,长，宽]，此时batch_size=1
img = img.to(device)
output = model(img)
pred = output.max(1, keepdim=True)[1]
pred = torch.squeeze(pred)
print('检测结果为：%d' % (pred.cpu().numpy()))四、模型序列化、可视化分析
基于上述工作，我们已经得到.pt格式的模型文件，想要在libtorch中加载模型，还要做的工作就是模型序列化，以下是python环境中序列化代码，注：需要读入一张图，最后将.pth格式文件转为.pt格式文件。
import torch
import cv2
import torch.nn.functional as F
from main import Net
from torchvision import datasets, transforms
from PIL import Image
if __name__ == '__main__':
device = torch.device('cpu') # 使用cpu进行推理
model = torch.load('pytorch_mnist.pth') # 加载模型
model = model.to(device)
model.eval() # 把模型转为test模式
img = cv2.imread("9.jpg", 0) # 读取要预测的灰度图片
img = Image.fromarray(img)
trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img = trans(img)
img = img.unsqueeze(0) # 图片扩展多一维,[batch_size,通道,长，宽]
img = img.to(device)
traced_net = torch.jit.trace(model, img)
traced_net.save("pytorch_mnist.pt")
print("模型序列化导出成功")生成的.pt格式文件将会部署到libtorch环境中使用，下一节详细讲解。这里最后我们利用链接：https://netron.app/的工具，将上述序列化后的模型文件(即：pytorch_mnist.pt)，进行可视化，如图（这里再将上述网络定义代码贴一遍，方便对比）：
class Net(nn.Module):
def __init__(self): # self指的是类实例对象本身(注意：不是类本身)。
# self不是关键词
# super 用于继承，https://www.runoob.com/python/python-func-super.html
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
# input:28*28
x = self.conv1(x) # -> (28 - 3 + 1 = 26),26*26*32
x = F.relu(x)
# input:26*26*32
x = self.conv2(x) # -> (26 - 3 + 1 = 24),24*24*64
# input:24*24*64
x = F.relu(x)
x = F.max_pool2d(x, 2)# -> 12*12*64 = 9216
x = self.dropout1(x) #不改变维度
x = torch.flatten(x, 1) # 9216*1
# w = 128*9216
x = self.fc1(x) # -> 128*1
x = F.relu(x)
x = self.dropout2(x)
# w = 10*128
x = self.fc2(x) # -> 10*1
output = F.log_softmax(x, dim=1) # softmax归一化
return output