devbox
寸_铁
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

模型训练+验证(train+demo)套路_一个模型的train文件

作者：寸_铁 | 2024-08-09 22:38:37

赞

踩

一个模型的train文件

以 CIFAR10 数据集为例，分类问题（10分类）

模型训练(train)套路

model.py

  1. import torch
  2. from torch import nn
  3.  
  4. # 搭建神经网络
  5. class Overwatch(nn.Module):
  6.     def __init__(self):
  7.         super(Overwatch, self).__init__()
  8.         self.model = nn.Sequential(
  9.             nn.Conv2d(3, 32, 5, 1, 2),
  10.             nn.MaxPool2d(2),
  11.             nn.Conv2d(32, 32, 5, 1, 2),
  12.             nn.MaxPool2d(2),
  13.             nn.Conv2d(32, 64, 5, 1, 2),
  14.             nn.MaxPool2d(2),
  15.             nn.Flatten(),
  16.             nn.Linear(64*4*4, 64),
  17.             nn.Linear(64, 10),
  18.         )
  19.  
  20.     def forward(self, x):
  21.         x = self.model(x)
  22.         return x
  23.  
  24. if __name__ == '__main__':  # 输入main即可跳出完整代码
  25.     OW = Overwatch()
  26.     input = torch.ones((64, 3, 32, 32))
  27.     output = OW(input)
  28.     print(output.shape)

train.py

  1. from model import *  # 引入模型
  2. import torchvision.datasets
  3. from torch.utils.data import DataLoader
  4.  
  5. # 准备数据集，CIFAR10数据集是PIL Image，要转换为tensor数据类型
  6. train_data = torchvision.datasets.CIFAR10('../data', train=True, transform=torchvision.transforms.ToTensor(),download=True)
  7. test_data = torchvision.datasets.CIFAR10('../data', train=False, transform=torchvision.transforms.ToTensor(),download=True)
  8.  
  9. # 查看训练数据集和测试数据集数量
  10. train_data_size = len(train_data)
  11. test_data_size = len(test_data)
  12.  
  13. print('训练数据集长度为：{}'.format(train_data_size))
  14. print('测试数据集长度为：{}'.format(test_data_size))
  15.  
  16. # 利用DataLoader加载数据集
  17. train_dataloader = DataLoader(train_data, batch_size=64)
  18. test_dataloader = DataLoader(test_data, batch_size=64)
  19.  
  20. # 创建网络模型
  21. OW = Overwatch()
  22.  
  23. # 损失函数
  24. loss_function = nn.CrossEntropyLoss()
  25.  
  26. # 优化器
  27. learning_rate = 0.01
  28. optimizer = torch.optim.SGD(OW.parameters(), lr=learning_rate)
  29.  
  30. #  设置训练网络的参数
  31. #  记录训练的次数
  32. total_train_step = 0
  33. #  记录测试的次数
  34. total_test_step = 0
  35. #  训练的轮数
  36. epoch = 10
  37.  
  38. for i in range(epoch):  # 0~9
  39.     print('-----第 {} 轮训练开始-----'.format(i+1))  # 1~10
  40.     # 训练步骤开始
  41.     for data in train_dataloader:
  42.         imgs, targets = data
  43.         outputs = OW(imgs)
  44.         loss = loss_function(outputs, targets)  # 输出与target计算loss
  45.  
  46.         # 优化器优化模型
  47.         optimizer.zero_grad()  # 梯度清零：把上一步训练的每个参数的梯度清零
  48.         loss.backward()  # 反向传播：调用反向传播得到每个要更新参数的梯度
  49.         optimizer.step()  # 参数优化：每个参数根据上一步得到的梯度进行优化
  50.  
  51.         total_train_step = total_train_step + 1
  52.         if total_train_step % 100 == 0:  # 每100次才打印
  53.             print('训练次数：{},Loss：{}'.format(total_train_step, loss.item()))

 tips:

print(loss) 和 print(loss.item()) 的区别

  1. import torch
  2. a = torch.tensor(5)
  3. print(a)              # tensor(5)
  4. print(a.item())       # 5

如何知道模型是否训练好，或达到需求？

每次训练完一轮就进行测试，以测试数据集上的损失或正确率来评估模型有没有训练好

测试过程中不需要对模型进行调优，利用现有模型进行测试

  1.     # 测试步骤开始（评估模型）
  2.     total_test_loss = 0
  3.     with torch.no_grad():  # 无需梯度调优
  4.         for data in test_dataloader:
  5.             imgs, targets = data
  6.             outputs = OW(imgs)
  7.             loss = loss_function(outputs, targets)
  8.             total_test_loss = total_test_loss + loss.item()
  9.     print('整体测试集上的Loss：{}'.format(total_test_loss))

与 TensorBoard 结合

正确率的实现（对分类问题） 

即便得到整体测试集上的 loss，也不能很好说明在测试集上的表现效果

  • 在分类问题中可以用正确率表示（下述代码改进）
  • 在目标检测/语义分割中，可以把输出放在tensorboard里显示，看测试结果

  1. outputs = torch.tensor([[0.1,0.2],
  2.                         [0.3,0.4]])
  3. print(outputs.argmax(1))  # 0或1表示方向，1为横向比较大小. 运行结果：tensor([1, 1])

argmax相当于就是返回数组中最大值的索引

 preds 预测

  1. preds = outputs.argmax(1)
  2. targets = torch.tensor([0,1])
  3. print(preds == targets)  # tensor([False,  True])
  4. print(preds == targets.sum())  # tensor(1)，对应位置相等的个数
  1.     # 测试步骤开始（评估模型）
  2.     # 初始化测试集上的总损失和总准确率
  3.     total_test_loss = 0
  4.    *total_accuracy = 0
  5.  
  6.     with torch.no_grad():  # 无需梯度调优
  7.         for data in test_dataloader:
  8.             imgs, targets = data
  9.             outputs = OW(imgs)
  10.             loss = loss_function(outputs, targets)
  11.             total_test_loss = total_test_loss + loss.item()
  12.            *accuracy = (outputs.argmax(1) == targets).sum()  # 得到预测正确的样本数量
  13.            *total_accuracy = total_accuracy + accuracy
  14.  
  15.     print('整体测试集上的Loss：{}'.format(total_test_loss))
  16.    *print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
  17.  
  18.     # 画test_accuracy & test_loss
  19.    *writer.add_scalar('test_accuracy', total_accuracy/test_data_size, total_test_step)
  20.     writer.add_scalar('test_loss', total_test_loss, total_test_step)
  21.     total_test_step = total_test_step + 1

完整代码

  1. from model import *  # 引入模型
  2. import torchvision.datasets
  3. from torch.utils.data import DataLoader
  4. from torch.utils.tensorboard import SummaryWriter
  5.  
  6. # 准备数据集，CIFAR10数据集是PIL Image，要转换为tensor数据类型
  7. train_data = torchvision.datasets.CIFAR10('../data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
  8. test_data = torchvision.datasets.CIFAR10('../data', train=False, transform=torchvision.transforms.ToTensor(), download=True)
  9.  
  10. # 查看训练数据集和测试数据集数量
  11. train_data_size = len(train_data)
  12. test_data_size = len(test_data)
  13.  
  14. print('训练数据集长度为：{}'.format(train_data_size))
  15. print('测试数据集长度为：{}'.format(test_data_size))
  16.  
  17. # 利用DataLoader加载数据集
  18. train_dataloader = DataLoader(train_data, batch_size=64)
  19. test_dataloader = DataLoader(test_data, batch_size=64)
  20.  
  21. # 创建网络模型
  22. OW = Overwatch()
  23.  
  24. # 损失函数
  25. loss_function = nn.CrossEntropyLoss()
  26.  
  27. # 优化器
  28. learning_rate = 0.01
  29. optimizer = torch.optim.SGD(OW.parameters(), lr=learning_rate)
  30.  
  31. # 设置训练网络的参数
  32. total_train_step = 0  # 记录训练的次数
  33. total_test_step = 0  # 记录测试的次数
  34. epoch = 10  # 训练的轮数
  35.  
  36. # 添加TensorBoard
  37. writer = SummaryWriter('../logs')
  38.  
  39. for i in range(epoch):
  40.     print('-----第 {} 轮训练开始-----'.format(i+1))
  41.     # 训练步骤开始
  42.     OW.train()
  43.     for data in train_dataloader:
  44.         imgs, targets = data
  45.         outputs = OW(imgs)
  46.         loss = loss_function(outputs, targets)  # 输出与target计算loss
  47.  
  48.         # 优化器优化模型
  49.         optimizer.zero_grad()  # 梯度清零：把上一步训练的每个参数的梯度清零
  50.         loss.backward()  # 反向传播：调用反向传播得到每个要更新参数的梯度
  51.         optimizer.step()  # 参数优化：每个参数根据上一步得到的梯度进行优化
  52.  
  53.         total_train_step = total_train_step + 1
  54.         if total_train_step % 100 == 0:  # 每100次才打印
  55.             print('训练次数：{},Loss：{}'.format(total_train_step, loss.item()))
  56.  
  57.             # 画train_loss
  58.             writer.add_scalar('train_loss', loss.item(), total_train_step)
  59.  
  60.     # 测试步骤开始（评估模型）
  61.     OW.eval()
  62.     # 初始化测试集上的总损失和总准确率
  63.     total_test_loss = 0
  64.     total_accuracy = 0
  65.     with torch.no_grad():  # 无需梯度调优
  66.         for data in test_dataloader:
  67.             imgs, targets = data
  68.             outputs = OW(imgs)
  69.             loss = loss_function(outputs, targets)
  70.             total_test_loss = total_test_loss + loss.item()
  71.             accuracy = (outputs.argmax(1) == targets).sum()  # 得到预测正确的样本数量
  72.             total_accuracy = total_accuracy + accuracy
  73.  
  74.     print('整体测试集上的Loss：{}'.format(total_test_loss))
  75.     print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
  76.  
  77.     # 画test_accuracy & test_loss
  78.     writer.add_scalar('test_accuracy', total_accuracy/test_data_size, total_test_step)
  79.     writer.add_scalar('test_loss', total_test_loss, total_test_step)
  80.     total_test_step = total_test_step + 1
  81.  
  82.     # 保存模型
  83.     torch.save(OW, 'OW_{}.pth'.format(i))  # 每一轮保存一个结果
  84.     # torch.save(OW.state_dict,'OW_{}.pth'.format(i))  save2
  85.     print('模型已保存')
  86.  
  87. writer.close()

总结：

准备数据→加载数据→准备模型→设置损失函数→设置优化器→开始训练→最后验证→结果聚合展示

利用GPU训练

方式一

网络模型： 

数据（包括输入、标注）：

训练过程

损失函数：

更好的写法：这种写法在CPU和GPU上都可以跑，优先在GPU上跑

比较CPU/GPU训练时间 

为了比较时间，引入 time 这个 package

 

方式二（更常用）


以下两种写法对于单显卡来说等价：

  1. device = torch.device("cuda")
  2. device = torch.device("cuda:0")

语法糖（一种语法的简写），程序在 CPU 或 GPU/cuda 环境下都能运行： 

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

通过这个变量可以控制是在CPU上运行还是GPU（改成 "cuda" 或 "cuda:0" ）上运行

  1. OW = OW.to(device)
  2. loss_function.to(device)
  3. imgs = imgs.to(device)
  4. targets = targets.to(device)
train_gpu.py
  1. from model import *  # 引入模型
  2. import torch
  3. import time
  4. import torchvision.datasets
  5. from torch import nn
  6. from torch.utils.data import DataLoader
  7. from torch.utils.tensorboard import SummaryWriter
  8.  
  9. # 定义训练设备
  10. device = torch.device('cuda')
  11.  
  12. # 准备数据集，CIFAR10数据集是PIL Image，要转换为tensor数据类型
  13. train_data = torchvision.datasets.CIFAR10('../data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
  14. test_data = torchvision.datasets.CIFAR10('../data', train=False, transform=torchvision.transforms.ToTensor(), download=True)
  15.  
  16. # 查看训练数据集和测试数据集数量
  17. train_data_size = len(train_data)
  18. test_data_size = len(test_data)
  19.  
  20. print('训练数据集长度为：{}'.format(train_data_size))
  21. print('测试数据集长度为：{}'.format(test_data_size))
  22.  
  23. # 利用DataLoader加载数据集
  24. train_dataloader = DataLoader(train_data, batch_size=64)
  25. test_dataloader = DataLoader(test_data, batch_size=64)
  26.  
  27. # 创建网络模型
  28. OW = Overwatch()
  29. # 将OW送给device
  30. OW = OW.to(device)
  31.  
  32. # 损失函数
  33. loss_function = nn.CrossEntropyLoss()
  34. loss_function.to(device)
  35.  
  36. # 优化器
  37. learning_rate = 0.01
  38. optimizer = torch.optim.SGD(OW.parameters(), lr=learning_rate)
  39.  
  40. #  设置训练网络的参数
  41. #  记录训练的次数
  42. total_train_step = 0
  43. #  记录测试的次数
  44. total_test_step = 0
  45. #  训练的轮数
  46. epoch = 10
  47.  
  48. # 添加TensorBoard
  49. writer = SummaryWriter('../logs')
  50. start_time = time.time()  # 记录下此时的时间，赋值给开始时间
  51. for i in range(epoch):
  52.     print('-----第 {} 轮训练开始-----'.format(i+1))
  53.     # 训练步骤开始
  54.     OW.train()
  55.     for data in train_dataloader:
  56.         imgs, targets = data
  57.         imgs = imgs.to(device)
  58.         targets = targets.to(device)
  59.         outputs = OW(imgs)
  60.         loss = loss_function(outputs, targets)  # 输出与target计算loss
  61.  
  62.         # 优化器优化模型
  63.         optimizer.zero_grad()  # 梯度清零：把上一步训练的每个参数的梯度清零
  64.         loss.backward()  # 反向传播：调用反向传播得到每个要更新参数的梯度
  65.         optimizer.step()  # 参数优化：每个参数根据上一步得到的梯度进行优化
  66.  
  67.         total_train_step = total_train_step + 1
  68.         if total_train_step % 100 == 0:  # 每100次才打印
  69.             end_time = time.time()
  70.             print(end_time - start_time)  # 第一次训练100次所花费的时间
  71.             print('训练次数：{},Loss：{}'.format(total_train_step, loss.item()))
  72.  
  73.             # 画train_loss
  74.             writer.add_scalar('train_loss', loss.item(), total_train_step)
  75.  
  76.     # 测试步骤开始（评估模型）
  77.     # 初始化测试集上的总损失和总准确率
  78.     OW.eval()
  79.     total_test_loss = 0
  80.     total_accuracy = 0
  81.     with torch.no_grad():  # 无需梯度调优
  82.         for data in test_dataloader:
  83.             imgs, targets = data
  84.             imgs = imgs.to(device)
  85.             targets = targets.to(device)
  86.             outputs = OW(imgs)
  87.             loss = loss_function(outputs, targets)
  88.             total_test_loss = total_test_loss + loss.item()
  89.             accuracy = (outputs.argmax(1) == targets).sum()  # 得到预测正确的样本数量
  90.             total_accuracy = total_accuracy + accuracy
  91.  
  92.     print('整体测试集上的Loss：{}'.format(total_test_loss))
  93.     print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
  94.  
  95.     # 画test_accuracy & test_loss
  96.     writer.add_scalar('test_accuracy', total_accuracy/test_data_size, total_test_step)
  97.     writer.add_scalar('test_loss', total_test_loss, total_test_step)
  98.     total_test_step = total_test_step + 1
  99.  
  100.     # 保存模型
  101.     torch.save(OW, 'OW_{}.pth'.format(i))  # 每一轮保存一个结果
  102.     #  torch.save(OW.state_dict,'OW_{}.pth'.format(i))  save2
  103.     print('模型已保存')
  104.  
  105. writer.close()


模型验证(demo)套路

核心：利用已经训练好的模型，给它提供输入进行测试

注意：load的模型当时保存的时候是用gpu训练的，因此这里需要定义device到cuda，再把image进行to(device)操作：device = torch.device('cuda')、image = image.to(device)

或者 model =torch.load("OW_30.pth",map_location=torch.device('cpu'))

采用GPU训练的模型，在CPU上加载，要从GPU上映射到CPU上

模型要求是四维的输入 [batch_size,channel,length,width]，但是获得的图片是三维的 —— 图片没有指定 batch_size（网络训练过程中是需要 batch_size 的，而图片输入是三维的，需要reshape() 一下）

test.py：

  1. import torch
  2. import torchvision.transforms
  3. from PIL import Image
  4. # from model import *
  5.  
  6. # device = torch.device('cuda')
  7. image_path = '../imgs/dog.png'  # 或绝对路径
  8. image = Image.open(image_path)
  9. print(image)
  10. image = image.convert('RGB')
  11. '''因为png格式是四通道，除了RGB三通道外，还有一个透明度通道，所以要调用上述语句保留其颜色通道
  12. 如果图片本来就是三颜色通道，经过此操作，不变。加上这一步后，可以适应 png jpg 各种格式的图片'''
  13.  
  14. # 网络模型的输入只能是32x32，进行一个Resize()
  15. transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),  # 32x32大小的PIL Image
  16.                                             torchvision.transforms.ToTensor()])  # 转为Tensor数据类型
  17. image = transform(image)
  18. print(image.shape)  # torch.Size([3, 32, 32])
  19.  
  20. # model = torch.load('OW_9.pth')
  21. model =torch.load("OW_9.pth",map_location=torch.device('cpu'))
  22. print(model)
  23. image = torch.reshape(image,(1, 3, 32, 32))
  24. # image = image.to(device)
  25. model.eval()  # 模型转化为测试类型
  26. with torch.no_grad():  # 节约内存和性能
  27.     output = model(image)
  28. print(output)
  29. print(output.argmax(1))

CIFAR10 对应的真实类别：

  

 

GitHub

下载代码用 PyCharm 打开后，查看代码里有没有 required=True，若有，删掉 required=True ，加一个默认值 default="./dataset/maps" ，就可以在 PyCharm 里右键运行了

即找到所有 required=True 的参数，将它删去并添加上默认值default 

 required的意思就是执行代码的时候需要手动输入文件目录，改为默认之后就不用再输入了

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/寸_铁/article/detail/955457
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号