CIFAR10数据集下载
CIFAR10数据集包含10个类别，图像尺寸为 3×32×32
在这里插入图片描述
官方下载地址很慢，这里给一个百度云 :
下载后在项目目录新建一个data目录解压进去
导入相关包
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time
import copy
MINI_BATCH = 8 # 数据集的图片数量很大，无法一次性加载所有数据，所以一次加载一个mini-batch的图片
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # GPU可用则使用GPU
使用torchvision加载并且归一化训练和测试数据集
CIFAR10数据集的输出是范围在[0,1]之间的PILImage，我们将它转换并归一化范围在[-1,1]之间的Tensor
# ToTensor(): 将ndarrray格式的图像转换为Tensor张量
# Normalize(std, mean): std：每个通道颜色平均值，这里的平均值为0.5，私人数据集自己计算；mean：每个通道颜色标准偏差，(原始数据 - std) / mean 得到归一化后的数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
数据加载器
# 训练数据加载
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=MINI_BATCH, shuffle=True, num_workers=4)
# 测试数据加载
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=4)
定义卷积神经网络
我们实现这样一个简单的神经网络：
在这里插入图片描述
这个网络具有两个卷积层，两个池化层，三个全连接层，输入数据为 3×32×32 的Tensorr数据，输出数据维度为 1*10 ，表示图片属于10个类别的概率，图中数据维度变化说明：
- 二维卷积层输出大小 out = (in - F + 2P) / S + 1 ，其中：
F： 卷积核大小 F×F
P： Padding，默认为0
S： 步长Stride，默认为1
如图中第一层卷积层 (32 - 5) / 1 + 1 = 28 - 池化层输出大小 ** out = (in - F) / S + 1** ，其中：
F： 池化窗口大小 F×F
S： 池化窗口移动的步长Stride，默认和池化窗口维度相同
如图中第二层池化层 (28 - 2) / 2 + 1 = 14
这部分可以写成一个独立的文件，在训练代码中引入此文件中的网络结构
# net.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5) # 卷积层：3通道到6通道，卷积5*5
self.conv2 = nn.Conv2d(6, 16, 5) # 卷积层：6通道到16通道，卷积5*5
self.pool = nn.MaxPool2d(2, 2) # 池化层，在2*2窗口上进行下采样
# 三个全连接层 ：16*5*5 -> 120 -> 84 -> 10
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)) # F.relu 是一个常用的激活函数
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.pool(x)
x = x.view(-1, 16 * 5 * 5) # 变换数据维度为 1*(16*5*5)，-1表示根据后面推测
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
定义一个通用的训练函数，每个epoch进行一次测试，最后返回具有最优参数的模型
def train(model, criterion, optimizer, epochs):
since = time.time()
best_acc = 0.0 # 记录模型测试时的最高准确率
best_model_wts = copy.deepcopy(model.state_dict()) # 记录模型测试出的最佳参数
for epoch in range(epochs):
print('-' * 30)
print('Epoch {}/{}'.format(epoch+1, epochs))
# 训练模型
running_loss = 0.0
for i, data in enumerate(trainloader):
inputs, labels = data
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
# 前向传播，计算损失
outputs = net(inputs)
loss = criterion(outputs, labels)
# 反向传播+优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
# 每1000批图片打印训练数据
if (i != 0) and (i % 1000 == 0):
print('step: {:d}, loss: {:.3f}'.format(i, running_loss/1000))
running_loss = 0.0
# 每个epoch以测试数据的整体准确率为标准测试一下模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
acc = correct / total
if acc > best_acc: # 当前准确率更高时更新
best_acc = acc
best_model_wts = copy.deepcopy(model.state_dict())
time_elapsed = time.time() - since
print('-' * 30)
print('训练用时： {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
print('最高准确率: {}%'.format(100 * best_acc))
# 返回测试出的最佳模型
model.load_state_dict(best_model_wts)
return model
接下来只需要定义好损失函数和优化器然后调用训练函数训练模型即可
from net import Net
net = Net()
net.to(DEVICE)
# 使用分类交叉熵 Cross-Entropy 作损失函数，动量SGD做优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 训练10个epoch
net = train(net, criterion, optimizer, 10)
# 保存模型参数
torch.save(net.state_dict(), 'net_dict.pt')
在这里插入图片描述
在这里插入图片描述
测试模型
# 图像类别
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = Net()
net.load_state_dict(torch.load('net_dict.pt')) # 加载各层参数
net.to(DEVICE)
# 整体正确率
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('整体准确率: {}%'.format(100 * correct / total))
print('=' * 30)
# 每一个类别的正确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images, labels = data
if torch.cuda.is_available():
images, labels = images.cuda(), labels.cuda()
outputs = net(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('{}的准确率 : {:.2f}%'.format(classes[i], 100 * class_correct[i] / class_total[i]))
在这里插入图片描述
下面看一下模型对测试集图片的一些预测结果
import matplotlib.pyplot as plt
import numpy as np
# 定义一个显示图片的函数
def imshow(img):
# 输入数据：torch.tensor[c, h, w]
img = img / 2 + 0.5 # 反归一
npimg = np.transpose(img.numpy(), (1, 2, 0)) # [c, h, w] -> [h, w, c]
plt.imshow(npimg)
plt.show()
# 取一批图片
testdata = iter(testloader)
images, labels = testdata.next()
imshow(torchvision.utils.make_grid(images))
print('真实类别: ', ' '.join('{}'.format(classes[labels[j]]) for j in range(4)))
# 预测是10个标签的权重，一个类别的权重越大，神经网络越认为它是这个类别，所以输出最高权重的标签。
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('预测结果: ', ' '.join('{}'.format(classes[predicted[j]]) for j in range(4)))
在这里插入图片描述
在这里插入图片描述
