模型训练步骤

模型训练步骤：

1. 准备数据集（train_dataset、test_dataset）

   import torchvision
   data = torchvision.datasets.数据集("数据集存放目录",train=True/False,
                                     transform=torchvision.transforms.ToTensor(),download=True)

2. 创建dataloader（train_dataloader、test_dataloader）

   from torch.utils.data import DataLoader
   train/test_dataloader = DataLoader(train_dataset,batch_size=64)

3. gpu or cpu

   device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

4. 搭建神经网络

   class ExamModuel(nn.Module):
       def __init__(self):
           super(ExamModuel, self).__init__()
           self.model = nn.Sequential(
               ...
           )
   
       def forward(self, x):
           x = self.model(x)
           return x

5. 创建网络模型 —> cuda

   ex_model = ExamModuel()
   ex_model.to(device)  

6. 损失函数 —> cuda

   loss_fn = nn.CrossEntropyLoss()
   loss_fn.to(device)

7. 优化器

   learning_rate = 1e-2    # 学习率
   optimizer = torch.optim.SGD(ex_model.parameters(),lr=learning_rate)

8. 一些参数

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

  
   writer = SummaryWriter("logs")

9. 开始训练

   迭代epoch次
   ex_model.train()
   for imgs, labels in dataloader:
   	imgs = imgs.to(device)
   	labels = labels.to(device)
   	outputs = ex_model(imgs)
   	loss = loss_fn(outputs, labels)
   	
   	# 优化器优化模型
   	optimizer.zero_grad()  # 梯度清零
       loss.backward()  # 反向传播,得到每一个参数节点的梯度
       optimizer.step()  # 进行优化
   	
   	# 结果可视化
   	print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
       writer.add_scalar("train_loss", loss.item(), total_train_step)

10. 开始测试

ex_model.eval()
with torch.no_grad():
       for imgs, labels in test_dataloader:
       img = img.to(device)
            target = target.to(device)
            output = ex_model(img)
            accuracy = (target == output.argmax(1)).sum()
            test_accuracy += accuracy
            loss = loss_fn(output, target)
            total_test_loss += loss.item()
            
            # acc、loss输出与可视化

11. 保存模型

torch.save(ex_model, "模型存放目录/ex_model_{}.pth".format(i))

12. 检验模型的效果（输入数据，得出输出结果）