完整的模型训练套路(以CIFAR10数据集)
(1)准备数据集
(2)dataloader加载数据集
(3)搭建网络模型
(4)创建网络模型实例
(5)定义损失函数
(6)定义优化器
(7)设置网络训练的参数
(8)开始训练
(9)验证模型
(10)最后保存模型
(11)将训练结果展示
(12)完整代码
一.准备数据集
#准备数据集
train_data=torchvision.datasets.CIFAR10("./data",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10("./data",train=False,transform=torchvision.transforms.ToTensor(),download=True)
二.DataLoader加载数据集
#利用DataLoader来加载数据集
train_data_loader=DataLoader(dataset=train_data,batch_size=64)
test_data_loader=DataLoader(test_data,batch_size=64)
三.搭建网络模型
class Tudui(nn.Module):
def __init__(self):
super().__init__()
self.seq1=nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=32,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32,out_channels=32,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32,out_channels=64,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self,x):
x=self.seq1(x)
return x
四.创建网络模型实例
#创建网络模型
tudui=Tudui()
五.创建损失函数
#创建损失函数
loss_fn=nn.CrossEntropyLoss()
六.创建优化器
#优化器
#1e-2=1 *(10)^-2=1/100
learning_rate=1e-2
optimizer=torch.optim.SGD(tudui.parameters(),lr=learning_rate)
七.设置网络训练参数
#设置训练网络的一些参数
#记录训练的次数
total_train_step=0
#记录测试次数
total_test_step=0
#训练的轮数
epoch=10
#添加tensorboard
writer=SummaryWriter("logs")
八.开始训练
for i in range(epoch):
print(f"---------------第{i+1}轮训练开始------------")
#训练步骤开始
for data in train_data_loader:
imgs,targets=data
outputs=tudui(imgs)
loss=loss_fn(outputs,targets)
#优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step+=1
if total_train_step%100==0:
print(f"训练次数:{total_train_step},loss:{loss.item()}")
writer.add_scalar("train_loss",loss.item(),total_train_step)
九.验证模型
#测试步骤开始
total_test_loss=0
with torch.no_grad():
for data in test_data_loader:
imgs,targets=data
outputs=tudui(imgs)
loss=loss_fn(outputs,targets)
total_test_loss+=loss.item()
print(f"整体测试集上的loss:{total_test_loss}")
writer.add_scalar("test_loss",total_test_loss,total_test_step)
total_test_step+=1
十.保存模型
torch.save(tudui,f"tudui_{i}.pth")
print("模型已保存")
十一.训练结果展示
writer.add_scalar("test_loss",total_test_loss,total_test_step)
十二.完整代码
import torch
import torchvision
from torch.utils.data import DataLoader
from torch import nn
from torch.utils.tensorboard import SummaryWriter
from model import *
#准备数据集
train_data=torchvision.datasets.CIFAR10("./data",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10("./data",train=False,transform=torchvision.transforms.ToTensor(),download=True)
train_data_size=len(train_data)
test_data_size=len(test_data)
#获得数据集的长度
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format((test_data_size)))
#利用DataLoader来加载数据集
train_data_loader=DataLoader(dataset=train_data,batch_size=64)
test_data_loader=DataLoader(test_data,batch_size=64)
#创建网络模型
tudui=Tudui()
#创建损失函数
loss_fn=nn.CrossEntropyLoss()
#优化器
#1e-2=1 *(10)^-2=1/100
learning_rate=1e-2
optimizer=torch.optim.SGD(tudui.parameters(),lr=learning_rate)
#设置训练网络的一些参数
#记录训练的次数
total_train_step=0
#记录测试次数
total_test_step=0
#训练的轮数
epoch=10
#添加tensorboard
writer=SummaryWriter("logs")
for i in range(epoch):
print(f"---------------第{i+1}轮训练开始------------")
#训练步骤开始
for data in train_data_loader:
imgs,targets=data
outputs=tudui(imgs)
loss=loss_fn(outputs,targets)
#优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step+=1
if total_train_step%100==0:
print(f"训练次数:{total_train_step},loss:{loss.item()}")
writer.add_scalar("train_loss",loss.item(),total_train_step)
#测试步骤开始
total_test_loss=0
with torch.no_grad():
for data in test_data_loader:
imgs,targets=data
outputs=tudui(imgs)
loss=loss_fn(outputs,targets)
total_test_loss+=loss.item()
print(f"整体测试集上的loss:{total_test_loss}")
writer.add_scalar("test_loss",total_test_loss,total_test_step)
total_test_step+=1
torch.save(tudui,f"tudui_{i}.pth")
print("模型已保存")
writer.close()
import torch
from torch import nn
#搭建神经网络
class Tudui(nn.Module):
def __init__(self):
super().__init__()
self.seq1=nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=32,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32,out_channels=32,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32,out_channels=64,kernel_size=5,stride=1,padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self,x):
x=self.seq1(x)
return x
if __name__=="__main__":
tudui=Tudui()
input=torch.ones((64,3,32,32))
output=tudui(input)
print(output.shape)