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

pytorch进阶学习（四）：使用不同分类模型进行数据训练（alexnet、resnet、vgg等）_pytorch训练分类模型

作者：Cpp五条 | 2024-06-07 17:51:32

赞

踩

pytorch训练分类模型

课程资源：5、帮各位写好了十多个分类模型，直接运行即可【小学生都会的Pytorch】_哔哩哔哩_bilibili

 

目录

一、项目介绍

 1. 数据集准备

2. 运行CreateDataset.py

3. 运行TrainModal.py 

4. 如何切换显卡型号

二、代码

1. CreateDataset.py

2.TrainModal.py 

3. 运行结果

一、项目介绍

 1. 数据集准备

数据集在data文件夹下

 

2. 运行CreateDataset.py

运行CreateDataset.py来生成train.txt和test.txt的数据集文件。

 

3. 运行TrainModal.py 

进行模型的训练，从torchvision中的models模块import了alexnet, vgg, resnet的多个网络模型，使用时直接取消注释掉响应的代码即可，比如我现在训练的是vgg11的网络。

  1.     # 不使用预训练参数
  2.    # model = alexnet(pretrained=False, num_classes=5).to(device) # 29.3%
  3.  
  4.     '''        VGG系列    '''
  5.     model = vgg11(weights=False, num_classes=5).to(device)   #  23.1%
  6.     # model = vgg13(weights=False, num_classes=5).to(device)   # 30.0%
  7.     # model = vgg16(weights=False, num_classes=5).to(device)
  8.  
  9.  
  10.     '''        ResNet系列    '''
  11.     # model = resnet18(weights=False, num_classes=5).to(device)    # 43.6%
  12.     # model = resnet34(weights=False, num_classes=5).to(device)
  13.     # model = resnet50(weights= False, num_classes=5).to(device)
  14.     #model = resnet101(weights=False, num_classes=5).to(device)   #  26.2%
  15.     # model = resnet152(weights=False, num_classes=5).to(device)

 同时需要注意的是， vgg11中的weights参数设置为false，我们进入到vgg的定义页发现weights即为是否使用预训练参数，设置为FALSE说明我们不使用预训练参数，因为vgg网络的预训练类别数为1000，而我们自己的数据集没有那么多类，因此不使用预训练。

 

把最后一行中产生的pth的文件名称改成对应网络的名称，如model_vgg11.pth。 

  1.     # 保存训练好的模型
  2.     torch.save(model.state_dict(), "model_vgg11.pth")
  3.     print("Saved PyTorch Model Success!")

4. 如何切换显卡型号

我在运行过程中遇到了“torch.cuda.OutOfMemoryError”的问题，显卡的显存不够，在服务器中使用查看显卡占用情况命令：

nvidia -smi

可以看到我一直用的是默认显卡0，使用情况已经到了100%，但是显卡1使用了67%，还用显存使用空间，所以使用以下代码来把显卡0换成显卡1.

  1. # 设置显卡型号为1
  2. import os
  3. os.environ['CUDA_VISIBLE_DEVICES'] = '1'

 

二、代码

1. CreateDataset.py

  1. '''
  2. 生成训练集和测试集，保存在txt文件中
  3. '''
  4. ##相当于模型的输入。后面做数据加载器dataload的时候从里面读他的数据
  5. import os
  6. import random#打乱数据用的
  7.  
  8. # 百分之60用来当训练集
  9. train_ratio = 0.6
  10.  
  11. # 用来当测试集
  12. test_ratio = 1-train_ratio
  13.  
  14. rootdata = r"data"  #数据的根目录
  15.  
  16. train_list, test_list = [],[]#读取里面每一类的类别
  17. data_list = []
  18.  
  19. #生产train.txt和test.txt
  20. class_flag = -1
  21. for a,b,c in os.walk(rootdata):
  22.     print(a)
  23.     for i in range(len(c)):
  24.         data_list.append(os.path.join(a,c[i]))
  25.  
  26.     for i in range(0,int(len(c)*train_ratio)):
  27.         train_data = os.path.join(a, c[i])+'\t'+str(class_flag)+'\n'
  28.         train_list.append(train_data)
  29.  
  30.     for i in range(int(len(c) * train_ratio),len(c)):
  31.         test_data = os.path.join(a, c[i]) + '\t' + str(class_flag)+'\n'
  32.         test_list.append(test_data)
  33.  
  34.     class_flag += 1
  35.  
  36. print(train_list)
  37. random.shuffle(train_list)#打乱次序
  38. random.shuffle(test_list)
  39.  
  40. with open('train.txt','w',encoding='UTF-8') as f:
  41.     for train_img in train_list:
  42.         f.write(str(train_img))
  43.  
  44. with open('test.txt','w',encoding='UTF-8') as f:
  45.     for test_img in test_list:
  46.         f.write(test_img)

2.TrainModal.py 

  1. '''
  2.     加载pytorch自带的模型，从头训练自己的数据
  3. '''
  4. import time
  5. import torch
  6. from torch import nn
  7. from torch.utils.data import DataLoader
  8. from utils import LoadData
  9.  
  10. import os
  11. os.environ['CUDA_VISIBLE_DEVICES'] = '1'
  12.  
  13.  
  14. from torchvision.models import alexnet  #最简单的模型
  15. from torchvision.models import vgg11, vgg13, vgg16, vgg19   # VGG系列
  16. from torchvision.models import resnet18, resnet34,resnet50, resnet101, resnet152    # ResNet系列
  17. from torchvision.models import inception_v3     # Inception 系列
  18.  
  19. # 定义训练函数，需要
  20. def train(dataloader, model, loss_fn, optimizer):
  21.     size = len(dataloader.dataset)
  22.     # 从数据加载器中读取batch（一次读取多少张，即批次数），X(图片数据)，y（图片真实标签）。
  23.     for batch, (X, y) in enumerate(dataloader):
  24.         # 将数据存到显卡
  25.         X, y = X.cuda(), y.cuda()
  26.  
  27.         # 得到预测的结果pred
  28.         pred = model(X)
  29.  
  30.         # 计算预测的误差
  31.         # print(pred,y)
  32.         loss = loss_fn(pred, y)
  33.  
  34.         # 反向传播，更新模型参数
  35.         optimizer.zero_grad()
  36.         loss.backward()
  37.         optimizer.step()
  38.  
  39.         # 每训练10次，输出一次当前信息
  40.         if batch % 10 == 0:
  41.             loss, current = loss.item(), batch * len(X)
  42.             print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
  43.  
  44.  
  45. def test(dataloader, model):
  46.     size = len(dataloader.dataset)
  47.     # 将模型转为验证模式
  48.     model.eval()
  49.     # 初始化test_loss 和 correct， 用来统计每次的误差
  50.     test_loss, correct = 0, 0
  51.     # 测试时模型参数不用更新，所以no_gard()
  52.     # 非训练， 推理期用到
  53.     with torch.no_grad():
  54.         # 加载数据加载器，得到里面的X（图片数据）和y(真实标签）
  55.         for X, y in dataloader:
  56.             # 将数据转到GPU
  57.             X, y = X.cuda(), y.cuda()
  58.             # 将图片传入到模型当中就，得到预测的值pred
  59.             pred = model(X)
  60.             # 计算预测值pred和真实值y的差距
  61.             test_loss += loss_fn(pred, y).item()
  62.             # 统计预测正确的个数
  63.             correct += (pred.argmax(1) == y).type(torch.float).sum().item()
  64.     test_loss /= size
  65.     correct /= size
  66.     print(f"correct = {correct}, Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
  67.  
  68.  
  69.  
  70.  
  71. if __name__=='__main__':
  72.     batch_size = 8
  73.  
  74.     # # 给训练集和测试集分别创建一个数据集加载器
  75.     train_data = LoadData("train.txt", True)
  76.     valid_data = LoadData("test.txt", False)
  77.  
  78.  
  79.     train_dataloader = DataLoader(dataset=train_data, num_workers=4, pin_memory=True, batch_size=batch_size, shuffle=True)
  80.     test_dataloader = DataLoader(dataset=valid_data, num_workers=4, pin_memory=True, batch_size=batch_size)
  81.  
  82.     # 如果显卡可用，则用显卡进行训练
  83.     device = "cuda" if torch.cuda.is_available() else "cpu"
  84.     print(f"Using {device} device")
  85.  
  86.  
  87.     '''
  88.         随着模型的加深，需要训练的模型参数量增加，相同的训练次数下模型训练准确率起来得更慢
  89.     '''
  90.  
  91.     # 不使用预训练参数
  92.    # model = alexnet(pretrained=False, num_classes=5).to(device) # 29.3%
  93.  
  94.     '''        VGG系列    '''
  95.     model = vgg11(weights=False, num_classes=5).to(device)   #  23.1%
  96.     # model = vgg13(weights=False, num_classes=5).to(device)   # 30.0%
  97.     # model = vgg16(weights=False, num_classes=5).to(device)
  98.  
  99.  
  100.     '''        ResNet系列    '''
  101.     # model = resnet18(weights=False, num_classes=5).to(device)    # 43.6%
  102.     # model = resnet34(weights=False, num_classes=5).to(device)
  103.     # model = resnet50(weights= False, num_classes=5).to(device)
  104.     #model = resnet101(weights=False, num_classes=5).to(device)   #  26.2%
  105.     # model = resnet152(weights=False, num_classes=5).to(device)
  106.  
  107.  
  108.  
  109.  
  110.     print(model)
  111.     # 定义损失函数，计算相差多少，交叉熵，
  112.     loss_fn = nn.CrossEntropyLoss()
  113.  
  114.     # 定义优化器，用来训练时候优化模型参数，随机梯度下降法
  115.     optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)  # 初始学习率
  116.  
  117.  
  118.     # 一共训练1次
  119.     epochs = 1
  120.     for t in range(epochs):
  121.         print(f"Epoch {t+1}\n-------------------------------")
  122.         time_start = time.time()
  123.         train(train_dataloader, model, loss_fn, optimizer)
  124.         time_end = time.time()
  125.         print(f"train time: {(time_end-time_start)}")
  126.         test(test_dataloader, model)
  127.     print("Done!")
  128.  
  129.     # 保存训练好的模型
  130.     torch.save(model.state_dict(), "model_vgg11.pth")
  131.     print("Saved PyTorch Model Success!")

3. 运行结果

vgg11的运行结果：，可以看到最后的准确率是24.8%，因为没有用预训练模型，所以准确率很低。

 

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

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

  

闽ICP备14008679号