-
-
导入适当的包
-
numpy ,torch, 等等
-
还可以导入 SummaryWriter 利用tensorboard来进行可视化
-
-
引入数据集
-
可以利用框架自带的datasets来进行数据导入。
-
也可以自己创建训练集和测试集
-
-
利用DataLoader来加载数据集(训练集和测试集)
-
DataLoader的作用是将数据集分块区域化,即将数据集划分为多个batch
-
-
搭建神经网络(通过类继承nn.Module)
-
类中实现 init 和 forward 方法
-
其中init方法是用来初始化神经网络,forward是自动执行神经网络输入,输出的过程,后续不用调用
-
-
定义损失函数和优化器
-
根据不同的问题来定义loss_fn ,回归一般用MSE或MAE分类一般用Cross-Entropy Loss
-
优化器包括随机梯度下降(SGD),动量优化器(Momentum)等
-
-
设置一下超参数即可以变的参数
-
训练次数,训练轮数等
-
-
训练开始
-
测试开始
-
保存模型
-
完整代码:
主函数:
#first dataset
import time
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import *
train_data=torchvision.datasets.CIFAR10(r'C:\Users\Felix\Desktop\learn_pytorch\dataset',train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10(r'C:\Users\Felix\Desktop\learn_pytorch\dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True)
#second length
train_size=len(train_data)
test_size=len(test_data)
print('训练集数据长度为',train_size)
print('测试集数据长度为',test_size)
#利用DataLoader来加载数据集
train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)
#搭建神经网络
tudui=Tudui()
tudui=tudui.cuda()
#损失函数
loss_fn=nn.CrossEntropyLoss()
loss_fn=loss_fn.cuda()
#优化器
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_train')
start_time=time.process_time()
for i in range(epoch):
print("-------第{}轮训练开始--------".format(i+1))
#训练步骤开始
tudui.train()#特殊的层才有
for data in train_dataloader:
imgs,targets=data
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))
writer.add_scalar("train_loss",loss.item(),total_train_step)
end_time=time.process_time()
print('运行时间:',end_time-start_time)
#测试步骤开始
tudui.eval()#特殊的层才有
total_test_loss=0
total_accuracy=0
with torch.no_grad():
for data in test_dataloader:
imgs,targets=data
imgs=imgs.cuda()
targets=targets.cuda()
outputs=tudui(imgs)
loss=loss_fn(outputs,targets)
total_test_loss+=loss.item()
accuracy=(outputs.argmax(1) == targets).sum()
total_accuracy+=accuracy.item()
print('整体测试集上的正确率',total_accuracy/test_size)
print("整体测试集上的loss:",total_test_loss)
writer.add_scalar('test_accuracy',total_accuracy,total_test_step)
writer.add_scalar("test_loss",total_test_loss,total_test_step+1)
total_test_step=total_test_step+1
torch.save(tudui,'tudui_{}.pth'.format(i))
print("模型已保存")
writer.close()
model
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(64*4*4,64),
nn.ReLU(),
nn.Linear(64,10),
nn.ReLU()
)
def forward(self,x):
x=self.model(x)
return x
if __name__ == '__main__':
tudui=Tudui()
imput=torch.ones((64,3,32,32))
output=tudui(imput)
print(output.shape)