完整的神经网络模型训练步骤——以CIFAR10数据集为例（pytorch版本）
完整步骤
准备数据集
DataLoader加载数据集
搭建网络模型
创建网络模型实例
定义损失函数
定义优化器
设置网络训练参数
开始训练
验证模型
保存模型
训练结果可视化（使用tensorboard完成）
数据集介绍
CIFAR10 是由 Hinton 的学生 Alex Krizhevsky、Ilya Sutskever 收集的一个用于普适物体识别的计算机视觉数据集，它包含 60000 张 32 X 32 的 RGB 彩色图片，总共 10 个分类。其中，包括 50000 张用于训练集，10000 张用于测试集。
CIFAR10数据集部分展示
CIFAR模型结构
CIFAR10模型结构
完整代码
CIFAR10网络模型结构：
# model.py文件
# 搭建CIFAR10神经网络
import torch
from torch import nn
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 验证搭建的网络是否正确
if __name__ == '__main__':
tudui = Tudui()
input = torch.ones((64, 3, 32, 32))
output = tudui(input)
print(output.shape)CPU版本训练代码
# train.py文件
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import *
# 准备数据集
# 训练数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=False,
transform=torchvision.transforms.ToTensor())
# 测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=False,
transform=torchvision.transforms.ToTensor())
# 观察训练数据集、测试数据集中的图像有多少张
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size)) # 训练数据集的长度为：50000
print("测试数据集的长度为：{}".format(test_data_size)) # 测试数据集的长度为：10000
# 使用dataloader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64) # 训练数据集加载
test_dataloader = DataLoader(test_data, batch_size=64) # 测试数据集加载
# 创建网络模型
tudui = Tudui()
# 创建损失函数
# 分类问题采用交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()
# 设定优化器
learning_rate = 0.01 # 0.01=1e-2
# learning_rate = 1e-2
optimizer = torch.optim.SGD(params=tudui.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter("./CIFAR10_train")
for i in range(epoch):
print("———————第{}轮训练开始——————".format(i+1))
# 训练步骤开始
for data in train_dataloader:
imgs, targets = data
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad() # 梯度清零
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤
total_test_loss = 0
total_accuracy_num = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy_num = total_accuracy_num + accuracy
print("整体测试集上的Loss：{}".format(total_test_loss))
print("整体测试集上的正确率：{}".format(total_accuracy_num/test_data_size))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
writer.add_scalar(("test_accuracy", (total_accuracy_num/test_data_size), total_test_step))
total_test_step = total_test_step + 1
# 保存每轮运行的模型
torch.save(tudui, "tudui_{}.pth".format(i))
print("模型已保存！")
writer.close()
GPU_1 版本
该方式是在CPU版本中的网络模型、数据、损失函数引入cuda即可。
例如：
对于网络模型：在CPU版本中，我们在实例化网络时只需 tudui = Tudui() 即可；而在GPU_1版本中，我们实例化网络时需要将 tudui = Tudui() 改写成 tudui = Tudui.cuda()。
对于数据：在CPU版本中，我们在训练、测试时所使用的数据只需
for data in train_dataloader: # 训练
imgs, targets = data
outputs = tudui(imgs)
for data in test_dataloader: # 测试
imgs, targets = data
outputs = tudui(imgs)
即可；而在GPU_1版本中，则为：
for data in train_dataloader: # 训练
imgs, targets = data
# 使用GPU
imgs = imgs.cuda()
targets = targets.cuda()
outputs = tudui(imgs)
for data in test_dataloader: # 测试
imgs, targets = data
# 使用GPU
imgs = imgs.cuda()
targets = targets.cuda()
outputs = tudui(imgs)
对于损失函数：在CPU版本中，我们所使用的损失函数只需
# 创建损失函数 # 分类问题采用交叉熵损失函数 loss_fn = nn.CrossEntropyLoss()
即可；而在GPU_1版本中，则有：
# 创建损失函数 # 分类问题采用交叉熵损失函数 loss_fn = nn.CrossEntropyLoss() loss_fn = loss_fn.cuda() # 使用GPU
GPU_1版本训练代码
# train_GPU_1.py文件
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
# 准备数据集
# 训练数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=False,
transform=torchvision.transforms.ToTensor())
# 测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=False,
transform=torchvision.transforms.ToTensor())
# 观察训练数据集、测试数据集中的图像有多少张
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size)) # 训练数据集的长度为：50000
print("测试数据集的长度为：{}".format(test_data_size)) # 测试数据集的长度为：10000
# 使用dataloader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64) # 训练数据集加载
test_dataloader = DataLoader(test_data, batch_size=64) # 测试数据集加载
# 创建网络模型
# 搭建CIFAR10神经网络
import torch
from torch import nn
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
tudui = Tudui()
tudui = tudui.cuda() # 使用GPU
# 创建损失函数
# 分类问题采用交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda() # 使用GPU
# 设定优化器
learning_rate = 0.01 # 0.01=1e-2
# learning_rate = 1e-2
optimizer = torch.optim.SGD(params=tudui.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter("./CIFAR10_train")
for i in range(epoch):
print("———————第{}轮训练开始——————".format(i+1))
# 训练步骤开始
for data in train_dataloader:
imgs, targets = data
# 使用GPU
imgs = imgs.cuda()
targets = targets.cuda()
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad() # 梯度清零
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤
total_test_loss = 0
total_accuracy_num = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
# 使用GPU
imgs = imgs.cuda()
targets = targets.cuda()
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy_num = total_accuracy_num + accuracy
print("整体测试集上的Loss：{}".format(total_test_loss))
test_accuracy = total_accuracy_num/test_data_size # 测试准确率
print("整体测试集上的正确率：{}".format(test_accuracy))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
writer.add_scalar("test_accuracy", test_accuracy, total_test_step)
total_test_step = total_test_step + 1
# 保存每轮运行的模型
torch.save(tudui, "tudui_{}.pth".format(i))
print("模型已保存！")
writer.close()
GPU_2 版本
该方式是在CPU版本中的网络模型、数据、损失函数使用to.(device)方法。
例如：
首先我们要定义设备，即device = torch.device("cpu")（cpu版本）或者device = torch.device("cuda")（GPU版本）。
对于网络模型：在CPU版本中，我们在实例化网络时只需 tudui = Tudui() 即可；而在GPU_2 版本中，我们实例化网络时需要将 tudui = Tudui() 改写成
tudui = Tudui() tudui.to(device)# 使用GPU
对于数据：在CPU版本中，我们在训练、测试时所使用的数据只需
for data in train_dataloader: # 训练
imgs, targets = data
outputs = tudui(imgs)
for data in test_dataloader: # 测试
imgs, targets = data
outputs = tudui(imgs)
即可；而在GPU_ 2 版本中，则为：
for data in train_dataloader: # 训练
imgs, targets = data
# 使用GPU
imgs = imgs.to(device)
targets = targets.to(device)
outputs = tudui(imgs)
for data in test_dataloader: # 测试
imgs, targets = data
# 使用GPU
imgs = imgs.to(device)
targets = targets.to(device)
outputs = tudui(imgs)
对于损失函数：在CPU版本中，我们所使用的损失函数只需
# 创建损失函数 # 分类问题采用交叉熵损失函数 loss_fn = nn.CrossEntropyLoss()
即可；而在GPU_2 版本中，则有：
# 创建损失函数 # 分类问题采用交叉熵损失函数 loss_fn = nn.CrossEntropyLoss() loss_fn.to(device) # 使用GPU
GPU_2版本训练代码
# train_GPU_2.py文件
import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
device = torch.device("cuda") # 定义训练设备
# 准备数据集
# 训练数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=False,
transform=torchvision.transforms.ToTensor())
# 测试数据集
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=False,
transform=torchvision.transforms.ToTensor())
# 观察训练数据集、测试数据集中的图像有多少张
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size)) # 训练数据集的长度为：50000
print("测试数据集的长度为：{}".format(test_data_size)) # 测试数据集的长度为：10000
# 使用dataloader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64) # 训练数据集加载
test_dataloader = DataLoader(test_data, batch_size=64) # 测试数据集加载
# 创建网络模型
# 搭建CIFAR10神经网络
import torch
from torch import nn
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
tudui = Tudui()
tudui.to(device)# 使用GPU
# 创建损失函数
# 分类问题采用交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device) # 使用GPU
# 设定优化器
learning_rate = 0.01 # 0.01=1e-2
# learning_rate = 1e-2
optimizer = torch.optim.SGD(params=tudui.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
# 训练轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter("./CIFAR10_train")
for i in range(epoch):
print("———————第{}轮训练开始——————".format(i+1))
# 训练步骤开始
for data in train_dataloader:
imgs, targets = data
# 使用GPU
imgs = imgs.to(device)
targets = targets.to(device)
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad() # 梯度清零
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤
total_test_loss = 0
total_accuracy_num = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
# 使用GPU
imgs = imgs.to(device)
targets = targets.to(device)
outputs = tudui(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy_num = total_accuracy_num + accuracy
print("整体测试集上的Loss：{}".format(total_test_loss))
test_accuracy = total_accuracy_num/test_data_size # 测试准确率
print("整体测试集上的正确率：{}".format(test_accuracy))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
writer.add_scalar("test_accuracy", test_accuracy, total_test_step)
total_test_step = total_test_step + 1
# 保存每轮运行的模型
torch.save(tudui, "tudui_{}.pth".format(i))
print("模型已保存！")
writer.close()
基于 tensorboard 的结果可视化
我们使用GPU（方式1，2均可）版本的代码，来完成tensorboard可视化训练过程及测试结果。
运行代码，会在项目文件夹中得到一个CIFAR10_train的文件夹。
这句代码的功劳！！ writer = SummaryWriter("./CIFAR10_train")
在 Terminal 中运行指令： tensorboard --logdir=CIFAR10_train，得到如下网址：http://localhost:6006/，打开该网址。
得到可视化结果，如下图所示。
tensorboard可视化界面
得到网络模型的训练损失
模型在训练集上的损失
模型在测试集上的损失
模型在测试集上的准确率
以上结果均为模型训练了10轮次所得结果！
预测实战
我们用训练好的模型来测试随机图片，看看效果如何。
我们使网络在训练数据集上训练30轮次，发现第30轮训练得到的模型较好，模型在测试集上的正确率达到64.56%，那么我们采用第30轮的模型用于验证随机图片。
每轮训练好模型（共30轮）
我们在网上随机选取一个 dog 的图片和 plane 的图片以及一个 cat 的照片，如下图所示。
dog
plane
cat
将这三张图片保存在项目文件夹的images文件夹中，如下图所示。
在pycharm中展示这三张图片。
十个类别真实标签。
验证dog的代码
import torch
import torchvision.transforms
from PIL import Image
from torch import nn
img_path = "./images/dog.png" # .表示当前目录
image = Image.open(img_path)
# print(image) #
image = image.convert("RGB") # 该代码的作用：
# 因为png格式的图片是四个通道，除了RGB三通道外，还有一个透明度通道。
# 所以我们调用image = image.convert("RGB")，保留其颜色通道。
# print(image) #
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
# print(image.size)
# CIFAR10网络模型结构
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 加载网络模型参数
model = torch.load("tudui_29.pth", map_location=torch.device("cpu"))
# map_location=torch.device("cpu") 将GPU版本的模型对应到CPU上
# print(model)
image = torch.reshape(image, (1, 3, 32, 32))
model.eval()
with torch.no_grad():
output = model(image)
print(output) # tensor([[ -3.3239, -5.6744, 5.5848, 3.1388, 4.8541, 13.1417, -3.0066, 4.5897, -14.3501, -7.2900]])
print(output.argmax(1)) # tensor([5]) 标签5对应的是狗
验证成功！
验证plane的代码
import torch
import torchvision.transforms
from PIL import Image
from torch import nn
img_path = "./images/plane.png" # .表示当前目录
image = Image.open(img_path)
# print(image) #
image = image.convert("RGB") # 该代码的作用：
# 因为png格式的图片是四个通道，除了RGB三通道外，还有一个透明度通道。
# 所以我们调用image = image.convert("RGB")，保留其颜色通道。
# print(image) #
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
# print(image.size)
# CIFAR10网络模型结构
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 加载网络模型参数
model = torch.load("tudui_29.pth", map_location=torch.device("cpu"))
# map_location=torch.device("cpu") 将GPU版本的模型对应到CPU上
# print(model)
image = torch.reshape(image, (1, 3, 32, 32))
model.eval()
with torch.no_grad():
output = model(image)
print(output) # tensor([[10.6781, -2.0109, 4.5609, 1.3858, -1.3748, -1.3599, 3.4883, -7.2311, -3.5088, -3.3030]])
print(output.argmax(1)) # tensor([0]) 标签0对应的是飞机
验证成功！
验证cat的代码
import torch
import torchvision.transforms
from PIL import Image
from torch import nn
img_path = "./images/cat.png" # .表示当前目录
image = Image.open(img_path)
# print(image) #
image = image.convert("RGB") # 该代码的作用：
# 因为png格式的图片是四个通道，除了RGB三通道外，还有一个透明度通道。
# 所以我们调用image = image.convert("RGB")，保留其颜色通道。
# print(image) #
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()])
image = transform(image)
# print(image.size)
# CIFAR10网络模型结构
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(), # 展平
nn.Linear(1024, 64), # 1024=64*4*4
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
# 加载网络模型参数
model = torch.load("tudui_29.pth", map_location=torch.device("cpu"))
# map_location=torch.device("cpu") 将GPU版本的模型对应到CPU上
# print(model)
image = torch.reshape(image, (1, 3, 32, 32))
model.eval()
with torch.no_grad():
output = model(image)
print(output) # tensor([[-1.4196, -5.0211, 1.2063, 4.8532, -2.0156, 2.7169, 0.4598, 0.5168, -0.3567, -0.8418]])
print(output.argmax(1)) # tensor([3]) 标签3对应的是猫
验证成功！
虽然训练30轮次的模型在测试集上的正确率只有64.56%，但我们经过随机验证可知，模型的效果还是很不错的。我们只验证了三张图片，当然我们也可以验证更多的图片。