模型训练套路：
- 1.准备数据集
- 2.加载数据集
- 3.搭建神经网络
- 4创建损失函数
- 5.优化器
- 6.设置训练网络的一些参数
- 7.添加tensorboard（方便观察）
- 8.开始训练
- .测试
- 9.保存神经网络
准备数据
#准备数据集
dataset_train=torchvision.datasets.CIFAR10("./dataset",train=True,transform=torchvision.transforms.ToTensor(),download=True)
dataset_test=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
加载数据
#加载数据集
loader_train=DataLoader(dataset_train,batch_size=16,drop_last=True,shuffle=False)
loader_test=DataLoader(dataset_test,batch_size=16,drop_last=True,shuffle=False)
搭建神经网络
class Net(nn.Module):
def __init__(self):
super().__init__()
self.model1=torch.nn.Sequential(
nn.Conv2d(3, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self,x):
return self.model1(x)
# #验证神经网络正确性
#
if __name__=="__main__":
net = Net()
input=torch.ones((64,3,32,32))
output=net(input)
print(output.shape)
创建损失函数
#创建损失函数
loss_cro=torch.nn.CrossEntropyLoss()
优化器
#创建优化器
optim=torch.optim.SGD(net.parameters(),lr=1e-3)
设置训练网络的一些参数
#设置训练网络的一些参数
total_train_stp=0#总训练次数
total_test_stp=0#总测试次数
epoch=20#训练轮次
添加tensorboard（方便观察）
#tensorboard
writer=SummaryWriter("./end")
开始训练
for i in range(epoch):
loss=0
print("-----------第{}轮训练开始--------------".format(i+1))
for data in loader_train:
imgs,targets=data
output=net(imgs)
loss=loss_cro(output,targets)
#优化器优化
optim.zero_grad()
loss.backward()
optim.step()
#记录训练次数
total_train_stp+=1
if total_train_stp%100==0:
print("训练次数：{}，loss:{}".format(total_train_stp,loss.item()))
writer.add_scalar("train_loss",loss.item(),total_train_stp)
测试
#测试
with torch.no_grad():
total_current_test=0
for data in loader_test:
imgs,targets=data
output=net(imgs)
loss=loss_cro(output,targets)
total_test_stp+=1
writer.add_scalar("test_loss",loss,total_test_stp)
accuracy=(output.argmax(1)==targets).sum()
total_current_test+=accuracy
print("整体测试的正确率为{}".format(total_current_test/(total_test_stp)*10))
writer.add_scalar("test_accuracy",total_current_test/(total_test_stp)*10,total_test_stp)
保存模型
torch.save(net,"./model3.pth")
总代码
总代码分在两个py文件中，分别是
- 负责训练模型的train.py
- 神经网络所在文件network.py
network.py
import torch
from torch import nn
class Net(nn.Module):
def __init__(self):
super().__init__()
self.model1=torch.nn.Sequential(
nn.Conv2d(3, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, padding=2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, padding=2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(1024, 64),
nn.Linear(64, 10)
)
def forward(self,x):
return self.model1(x)
# #验证神经网络正确性
#
if __name__=="__main__":
net = Net()
input=torch.ones((64,3,32,32))
output=net(input)
print(output.shape)
train.py
import torchvision
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from network import Net
#准备数据集
dataset_train=torchvision.datasets.CIFAR10("./dataset",train=True,transform=torchvision.transforms.ToTensor(),download=True)
dataset_test=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
#加载数据集
loader_train=DataLoader(dataset_train,batch_size=16,drop_last=True,shuffle=False)
loader_test=DataLoader(dataset_test,batch_size=16,drop_last=True,shuffle=False)
#搭建神经网络
net=Net()
#创建损失函数
loss_cro=torch.nn.CrossEntropyLoss()
#创建优化器
optim=torch.optim.SGD(net.parameters(),lr=1e-3)
#设置训练网络的一些参数
total_train_stp=0#总训练次数
total_test_stp=0#总测试次数
epoch=20#训练轮次
#tensorboard
writer=SummaryWriter("./end")
#开始训练
for i in range(epoch):
loss=0
print("-----------第{}轮训练开始--------------".format(i+1))
for data in loader_train:
imgs,targets=data
output=net(imgs)
loss=loss_cro(output,targets)
#优化器优化
optim.zero_grad()
loss.backward()
optim.step()
#记录训练次数
total_train_stp+=1
if total_train_stp%100==0:
print("训练次数：{}，loss:{}".format(total_train_stp,loss.item()))
writer.add_scalar("train_loss",loss.item(),total_train_stp)
#测试
with torch.no_grad():
total_current_test=0
for data in loader_test:
imgs,targets=data
output=net(imgs)
loss=loss_cro(output,targets)
total_test_stp+=1
writer.add_scalar("test_loss",loss,total_test_stp)
accuracy=(output.argmax(1)==targets).sum()
total_current_test+=accuracy
print("整体测试的正确率为{}".format(total_current_test/(total_test_stp)*10))
writer.add_scalar("test_accuracy",total_current_test/(total_test_stp)*10,total_test_stp)
torch.save(net,"./model3.pth")
print("模型已经保存")
writer.close()
最终运行截图
可以看出再训练了20轮后可以达到93%的正确率
