Search K
Appearance
Appearance
以CIFAR10数据集为例子,这个数据集有10个类别,对我们来说是个十分类问题 介绍:我们以这个例子为例,来介绍训练网络的套路
python
import torchvision
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="../data", train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10(root="../data", train=False, transform=torchvision.transforms.ToTensor(),
download=True)
# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
# 如果train_data_size=10, 训练数据集的长度为: 10
print("训练数据集的长度为: {}".format(train_data_size))
print("测试数据集的长度为: {}".format(test_data_size))
from torch.utils.data import DataLoader
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
CIFAR10其中有10个类别,所以这个网络应该是10分类的网络。
# model.py
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.model = nn.Sequential(
# 第一个卷积层: 输入3通道,输出32通道,kernel_size=5,stride=1,padding=2
nn.Conv2d(3, 32, 5, 1, 2),
# 最大池化层: kernel_size=2,将特征图尺寸减半
nn.MaxPool2d(2),
# 第二个卷积层: 输入32通道,输出32通道
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
# 第三个卷积层: 输入32通道,输出64通道
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
# 展平层: 将特征图转换为一维向量
nn.Flatten(),
# 第一个全连接层: 输入64*4*4,输出64
nn.Linear(64*4*4, 64),
# 输出层: 输入64,输出10个类别
nn.Linear(64, 10)
)
def forward(self, x):
# 前向传播
x = self.model(x)
return x
if __name__ == '__main__':
model = Model()
input = torch.ones(64, 3, 32, 32) # 64张图片,3通道,32x32的图片
output = model(input)
print(output.shape) # 输出:torch.Size([64, 10])
# 它返回64个数据,每个数据有10个数据,这10代表每张图片在这10个类别中的概率
nn.Sequential(序列模块): 是 PyTorch 中的一个模块,主要用来构建神经网络的顺序容器。它可以按给定的顺序将多个神经网络模块(如全连接层、卷积层、激活函数等)组合在一起,形成一个有序的网络结构,同时,在每个模块的内部,模块的forward方法会被自动执行。
from model import *
# 创建网络模型
model = Model()
# 损失函数
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失函数
为什么要用交叉熵损失函数? 因为是分类问题,所以损失函数,我们可以用交叉熵损失函数,交叉熵损失函数是用来衡量两个概率分布之间的距离,对于分类问题,我们可以用交叉熵损失函数来衡量模型输出的概率分布和真实标签的概率分布之间的距离。
torch.optim
里面有很多优化器,比如SGD,Adam等,这里我们选择SGD优化器(随机梯度下降),里面的参数定义好,需要优化哪些参数。
# 优化器
leraning_rate = 0.01 # 学习率,写在外面是方便查找
optimizer = torch.optim.SGD(model.parameters(), lr=leraning_rate)
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#设置训练的轮数
epoch=10
#开始训练
for i in range(epoch): # 训练10轮,训练的内容写在下面
print("----第{}轮训练开始----".format(i+1))
# 训练步骤开始
for data in train_dataloader: # 从训练集取数据
imgs, targets = data
outputs = model(imgs)
loss = loss_fn(outputs, targets) # 看看误差是多少,将输出 和 真实值放进去,然后得到了一个loss
# 接下来进行优化了(优化器优化模型)
optimizer.zero_grad() # 首先梯度清零
loss.backward() # 梯度清零的时候,我们调用我们求出来的这个损失, 反向传播,得到每个参数 节点的梯度
optimizer.step() # 再调用这个step, 对每个参数进行一个优化,更新参数
# 到这里,一个训练步骤就结束了
total_train_step += 1
print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))
egg: 这里的loss.item(),这里的item(),是可以将tensor转换为数值的。
问题来了:我们训练的时候,怎么知道模型训练的怎么样了,有没有训练好
所以,我们每训练一轮,就要在测试数据集上跑一遍,以测试数据集的损失、或者准确率,来评估这个模型有没有训练好。
#记录训练的次数
total_train_step=0
#记录测试的次数
total_test_step=0
#设置训练的轮数
epoch=10
# 添加tensorboard
writer = SummaryWriter("../logs")
#开始训练
for i in range(epoch): # 训练10轮,训练的内容写在下面
print("----第{}轮训练开始----".format(i+1))
# 训练步骤开始
for data in train_dataloader: # 从训练集取数据
imgs, targets = data
outputs = model(imgs)
loss = loss_fn(outputs, targets) # 看看误差是多少,将输出 和 真实值放进去,然后得到了一个loss
# 接下来进行优化了(优化器优化模型)
optimizer.zero_grad() # 首先梯度清零
loss.backward() # 梯度清零的时候,我们调用我们求出来的这个损失, 反向传播,得到每个参数 节点的梯度
optimizer.step() # 再调用这个step, 对每个参数进行一个优化,更新参数
# 到这里,一个训练步骤就结束了
total_train_step += 1
# print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))
if total_train_step % 100 == 0: # 每训练100次,打印一次
print("训练次数: {}, Loss: {}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
# 测试步骤开始
# 我们一般求整体测试集上的损失,所以,我们初始化一个total_test_loss,然后,在测试的时候,将每个测试数据集的损失加起来,得到整体测试集的损失
total_test_loss = 0 # 初始化测试损失
with torch.no_grad(): # 从名字就可以看出,测试的时候,不需要计算梯度,不用调优
for data in test_dataloader: # 从测试集取数据
imgs, targets = data
outputs = model(imgs)
loss = loss_fn(outputs, targets) # 计算损失
total_test_loss += loss.item() # 将每个测试数据集的损失加起来,得到整体测试集的损失
# print("测试次数: {}, Loss: {}".format(total_test_step, loss.item()))
print("整体测试集上的Loss: {}".format(total_test_loss))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
total_test_step += 1 # 测试次数加1
# 保存模型
torch.save(model, "model_{}.pth".format(i))
print("模型已保存")
writer.close()