Pytorch完整模型训练套路

Pytorch完整训练模型套路

• 以CIFAR_10模型为例

• 使用tensorboard将训练过程可视化

• 模型文件：model.py

  import torch
  from torch import nn
  
  # 搭建神经网络
  class CIFAR_10_Model(nn.Module):
      def __init__(self):
          super().__init__()
          self.model=nn.Sequential(
              nn.Conv2d(3,32,5,1,2),
              nn.MaxPool2d(2),
              nn.Conv2d(32,16,5,1,2),
              nn.MaxPool2d(2),
              nn.Flatten(),
              nn.Linear(64*4*4,64),
              nn.Linear(64,10)
  
          )
      def forward(self,input):
          return self.model(input)
  
  if __name__=="__main__":
      cifar10=CIFAR_10_Model()
      input=torch.ones((64,3,32,32))
      output=cifar10(input)
      print(output.shape)
  

• 训练文件train.py

  import torchvision
  from torch.utils.tensorboard import SummaryWriter
  from model import *
  from torch import nn
  from torch.utils.data import DataLoader
  
  # 准备数据集
  train_data=torchvision.datasets.CIFAR10(root="../data/dataset",train=True,
                                          transform=torchvision.transforms.ToTensor(),
                                          download=True)
  test_data=torchvision.datasets.CIFAR10(root="../data/dataset",train=False,
                                          transform=torchvision.transforms.ToTensor(),
                                          download=True)
  
  # 数据集的长度
  train_data_size=len(train_data)
  test_data_size=len(test_data)
  print(f"train_set长度：{train_data_size}")
  print(f"test_set长度：{test_data_size}")
  
  # 利用DataLoader加载数据
  train_dataloader=DataLoader(train_data,batch_size=64)
  test_dataloader=DataLoader(test_data,batch_size=64)
  
  # 创建CIFAR_10模型实例
  Cifar10=CIFAR_10_Model()
  
  #损失函数
  loss_fn=nn.CrossEntropyLoss()
  
  #学习率
  learning_rate=1e-2
  #优化器
  optimizer=torch.optim.SGD(Cifar10.parameters(),lr=learning_rate)
  
  #训练次数
  total_train_step=0
  #测试次数
  total_test_step=0
  #训练轮数
  epoch=10
  
  # 添加tensorboard
  writer=SummaryWriter("./dataset_logs")
  
  for i in range(epoch):
      print(f"----------第{i+1}轮训练开始----------")
  
      #开始训练
      Cifar10.train()
      for data in train_dataloader:
          imgs,targets=data
          outputs=Cifar10(imgs)
          loss=loss_fn(outputs,targets)
  
          #反向传播，梯度下降
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          total_train_step=total_test_step+1
          if (total_train_step+1)%100==0:
              print(f"训练次数：{total_train_step}，Loss：{loss}")
              writer.add_scalar("train_loss",loss.item(),total_train_step)
  
      #开始测试
      Cifar10.eval()
      total_test_loss=0
      total_accuray=0
      with torch.no_grad():
          for data in test_dataloader:
              imgs,targets=data
              outputs=Cifar10(imgs)
              loss=loss_fn(outputs,targets)
              total_test_loss=total_test_loss+loss.item()
              accuracy=(outputs.argmax(1)==targets).sum()
              total_accuray=total_accuray+accuracy
      print(f"整体测试集的Loss：{total_test_loss}")
      print(f"整体测试集的正确率：{total_accuray/test_data_size}")
      writer.add_scalar("test_loss",total_test_loss,total_test_step)
      writer.add_scalar("test_accuracy",total_accuray/test_data_size,total_test_step)
      total_test_step=total_test_step+1
  
  writer.close()
  
  

• 参考