参考:
一、实现LeNet网络模型
网络构造图
代码说明
- 此代码我们只需要理解初始化代码(
__init__(self)
) 和正向传播(forward(self,x)
)函数即可。 - 初始化代码用来初始化,正向传播用来定义数据流的传输过程
- 代码中的初始化部分中的conv1,conv2…对应上图中的用红色标的相同过程
'''
LeNet的实现
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
#定义第一层卷积层,RGB为三通道图,所以输入通道为3,输出为6,卷积核为5*5
#按照公式:输出规则=(输入规格-卷积核+2*padding)/步长 + 1,所以这里输出规格为28*28
#此函数详细信息详阅https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html?highlight=conv2d#torch.nn.Conv2d
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool1 = nn.MaxPool2d(2,2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(16 * 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)) #output shape: torch.Size([32, 6, 28, 28])
x = self.pool1(x) #output shape: torch.Size([32, 6, 14, 14])
x = F.relu(self.conv2(x)) #output shape: torch.Size([32, 16, 10, 10])
x = self.pool2(x) #output shape: torch.Size([32, 16, 5, 5])
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
input1 = torch.rand([32,3,32,32])
model = LeNet()
print(model)
output = model(input1)
- 代码中传输的四维数据表示 [batch数,通道数,高,宽],batch数即一次处理多少张图片
- 上述代码完美的对应了图中的网络结构
二、训练模型
载入训练集和验证集,并简单的看里面的图片数据
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
#显示一张图片
def imshow(img):
img = img / 2 + 0.5 #图片之前减0.5再除以0.5,我们现在进行反向操作进行还原
npimg = img.numpy() #转化为numpy格式
plt.imshow(np.transpose(npimg, (1, 2, 0))) #将[C,H,W]改成[H,W,C]的格式
plt.show()
#定义图像数据处理方式 先转Tensor再减均值除方差
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 4
#定义训练集
trainset = torchvision.datasets.CIFAR10(root='./dataset', train=False, #下载完成这里改为False
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=0)
#定义测试集
testset = torchvision.datasets.CIFAR10(root='./dataset', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
dataiter = iter(trainloader) #将trainloader转换成可迭代的方式,也可以将trainloader改成testloader来看测试集
images, labels = dataiter.next()
#输出图片对应的标签
print(' '.join('%5s' % classes[labels[j]] for j in range(batch_size)))
#显示图片
imshow(torchvision.utils.make_grid(images))
这里将测试集的图片进行显示出来(如下图,标签在最后),由于分辨率不高,因此显示的比较模糊。
![在这里插入图片描述](https://www.icode9.com/i/ll/?i=20210716203336957.png#pic_center
训练模型并保存
#定义损失函数有优化器
net = LeNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
#开始训练
for epoch in range(10): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
#获取图片数据和标签
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 500 == 499: # print every 2000 mini-batches
with torch.no_grad(): #测试,不进行反向传播
accuracy = 0.0
for ii,data in enumerate(testloader): #循环载入所有测试集图片
inputs,label = data
outputs = net(inputs)
predict = torch.max(outputs,dim=1)[1]
accuracy += (predict == label).sum().item()
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, i + 1, running_loss / 500,accuracy/len(testset)))
running_loss = 0.0
print('Finished Training')
torch.save(net.state_dict(),"LeNet/LeNet.pth") #保存模型参数
运行上述代码,最后一个epoch的结果如下
可以看到图片的正确率已经到了68%,同时在LeNet目录下产生了模型文件LeNet.pth。接下来看如何使用
三、预测图片
在这一部分,加载前面训练好的模型LeNet.pth
,在网上下载一张飞机,车,猫
等我们在classes
中定义过的类别的图片放到LeNet目录下,运行下面的代码(我使用的是一张马的图片)
import torch
import torchvision.transforms as transforms
from PIL import Image
from torchvision.transforms.transforms import Resize
from LeNetModel import LeNet
#必须采用和训练集相同的图片预处理方式
transform = transforms.Compose(
[transforms.Resize((32,32)), #模型输入规格是32*32,因此我们必须先把图片转成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('LeNet/LeNet.pth'))
im = Image.open("Lenet/horse.jpg")
#transform出来的矩阵是三维的,是[C,H,W]的格式
#必须写成[N,C,H,W]的形式
im = torch.unsqueeze(transform(im),dim = 0)
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs,dim = 1)[1].item()
print(classes[predict])
得到的结果是horse
,所以还是挺准的