25. 完整的训练套路(一)数据准备、模型搭建、模型训练

完整的训练套路(一): 数据准备、模型搭建、模型训练

1. 数据的准备

• 训练模型之前首先需要加载好数据，本文以CIFAR10中的数据为例构建训练集和测试集

  # load data
  train_data = torchvision.datasets.CIFAR10(root = "./data", train = True, 
                                         transform=torchvision.transforms.ToTensor())
  test_data = torchvision.datasets.CIFAR10(root = "./data", train = False, 
                                           transform=torchvision.transforms.ToTensor())
  

  • 上述代码下载了CIFAR10中的测试集和训练集并转换为了 Tensor 数据类型

• 对于训练集和测试集，使用DataLoader对数据集进行重新分组、封装从而更适合模型的训练

  # DataLoader load dataset
  train_dataloader = DataLoader(train_data, batch_size=64)
  test_dataloader = DataLoader(test_data, batch_size=64)
  

2. 模型的搭建

• 沿用前面的文章中针对CIFAR10构建的模型架构

  # create Model
  class Model(nn.Module):
      def __init__(self):
          super(Model, 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(1024, 64),
              nn.Linear(64, 10)
          )        
  
      def forward(self, x):
          x = self.model(x)
          return x
  model = Model()
  

• 在一个较为规范的模型搭建中，都会将模型单独存放在一个文件中。对于构建好的模型，最好是进行模型的验证，保证模型在后续的训练过程中不会因为维度问题而报错

  input = torch.ones(64, 3, 32, 32)
  output = model(input)
  print(output.size())
  

  • 上述代码创建了一个 64*3*32*32 的单位矩阵，与CIFAR10的数据维度相同，放入模型进行前向传播，确保模型没有逻辑问题

3. 模型的训练

• 模型训练之前，首先要创建具体的损失函数和要使用的优化器

  # create Loss Function
  loss_fn = nn.CrossEntropyLoss()
  
  # create optim
  learning_rate = 0.01
  optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
  

  • 针对分类问题，这里使用交叉熵损失函数；优化器使用 SGD 优化器
  • 注意优化器的两个参数，params = model.parameters() 以及 lr = learning_rate

• 创建一个模型训练过程中的计数器

  total_train_step = 0  # 训练次数
  total_test_step = 0 # 测试次数
  epoch = 10 # 训练轮数
  

• epoch设置了模型的训练轮数，每一轮都是去遍历所有的训练数据集，训练过程主要涉及以下步骤: (1)前向传播得输出； (2)损失函数求误差； (3)梯度置零后反向传播; (4)优化器优化模型参数

  for i in range(epoch):
      print(f"----------{i+1} 轮训练开始----------")
      
      # start train
      for data in train_dataloader:
          imgs, targets = data
          outputs = model(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()}")
  

  • 上述代码中的loss 和 loss.item()的区别是，后者是前者的具体的数值，举例说明如下

    a = torch.tensor([1])  # output:  tensor([1])
    a.item() # output: 1
    

• 在每一轮的训练结束后，我们往往要对测试集进行测试，观察整个模型在每一轮训练后的性能表现。

  # test model
  total_test_loss = 0
  with torch.no_grad(): # 不进行梯度变化
      for data in test_dataloader:
          imgs, targets = data
          outputs = model(imgs)
          loss = loss_fn(outputs, targets)  # 当前batch的误差
          total_test_loss += loss.item() # 计算整体数据集上的loss
  total_test_step += 1 # 测试次数加一
  print(f"整体数据集上的loss {total_test_loss}")
  

  • with torch.no_grad(): 这个很重要，在对测试集进行预测的过程中，同样会使用模型架构的前向传播，我们要设置整个测试集的预测过程中不进行梯度的变化！！保证模型的参数就是当前训练轮次得到的参数
  • total_test_loss:这是因为，每次遍历 test_dataloader 得到的loss只是一个batch_size中的测试集的损失，我们应该要获取的是整体测试集的损失！！

• 训练+测试部分的完整代码如下：

  for i in range(epoch):
      print(f"----------{i+1} 轮训练开始----------")
      # train model
      for data in train_dataloader:
          imgs, targets = data
          outputs = model(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()}")
      
      # test model
      total_test_loss = 0
      with torch.no_grad(): # 不进行梯度变化
          for data in test_dataloader:
              imgs, targets = data
              outputs = model(imgs)
              loss = loss_fn(outputs, targets)  # 当前batch的误差
              total_test_loss += loss.item() # 统计整体数据集上的loss
      total_test_step += 1 # 测试次数加一
      print(f"整体数据集上的loss {total_test_loss}")
  

  • 针对上述代码的一个样例输出如下：

    ----------1 轮训练开始----------
    训练次数 1400， loss 1.949365258216858
    训练次数 1500， loss 1.790314793586731
    训练次数 1600， loss 1.585450530052185
    训练次数 1700， loss 1.6580116748809814
    训练次数 1800， loss 1.830728530883789
    训练次数 1900， loss 1.6263401508331299
    训练次数 2000， loss 1.7490088939666748
    训练次数 2100， loss 1.45167076587677
    整体数据集上的loss 269.71709048748016