• AI项目十七：ResNet50训练部署教程

    若该文为原创文章，转载请注明原文出处。

    ResNet50训练主要还是想部署到RK3568开发板上，先记录下训练和转成ONNX模型过程。

    一、 Resnet50简介

           ResNet50网络是2015年由微软实验室的何恺明提出，获得ILSVRC2015图像分类竞赛第一名。在ResNet网络提出之前，传统的卷积神经网络都是将一系列的卷积层和池化层堆叠得到的，但当网络堆叠到一定深度时，就会出现退化问题。 残差网络的特点是容易优化，并且能够通过增加相当的深度来提高准确率。其内部的残差块使用了跳跃连接，缓解了在深度神经网络中增加深度带来的梯度消失问题。

    二、数据集下载

           本教程以车辆分类算法为例，数据集的百度网盘下载链接为:

    https://pan.baidu.com/s/1pkYm9AA3s3WDM7GecShlbQ 提取码：6666​

           解压完成后得到以下两个文件夹：

           打开可以看到一共10类汽车：

    三、环境搭建

    1、创建虚拟环境

    conda create -n Resnet50_env python=3.8 -y

    2、激活环境

    conda activate Resnet50_env

    注意：使用的是CPU版本，电脑无GPU

    3、安装环境

    1. pip install numpy
    2. pip install torch
    3. pip install torchvision
    4. pip install matplotlib

    至此，环境安装完成，开始训练

    四、 ResNet50图像分类训练

    直接上源码:train.py

    1. # -#-coding:utf-8 -*-
    2.  
    3. import os
    4. import numpy as np
    5. import torch
    6. import torch.nn as nn
    7. import torch.nn.functional as F
    8. import torch.optim as optim
    9. import torchvision
    10. from torch.autograd.variable import Variable
    11. from torch.utils.data import DataLoader
    12. from torchvision import datasets, transforms
    13. import matplotlib.pyplot as plt
    14. from PIL import ImageFile
    15. ImageFile.LOAD_TRUNCATED_IMAGES = True
    16.  
    17. # 2.定义超参数
    18. BATCH_SIZE = 16  # 每批处理的数据
    19. DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # 放在cuda或者cpu上训练
    20. EPOCHS = 15  # 训练数据集的轮次
    21. modellr = 1e-3
    22.  
    23. # 3.构建pipeline，对图像做处理
    24. pipeline = transforms.Compose([
    25.     # 分辨率重置为256
    26.     transforms.Resize(256),
    27.     # 对加载的图像作归一化处理， 并裁剪为[224x224x3]大小的图像(因为这图片像素不一致直接统一)
    28.     transforms.CenterCrop(224),
    29.     # 将图片转成tensor
    30.     transforms.ToTensor(),
    31.     # 正则化，模型出现过拟合现象时，降低模型复杂度
    32.     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    33. ])
    34.  
    35. # 图片路径(训练图片和测试图片的)
    36. base_dir_train = 'G:/enpei_Project_Code/22_Resnet50_bus/1.data/datasets/train'
    37. base_dir_val = 'G:/enpei_Project_Code/22_Resnet50_bus/1.data/datasets/val'
    38.  
    39. # 4. 加载数据集
    40. train_dataset = datasets.ImageFolder(root=base_dir_train, transform=pipeline)
    41. print("train_dataset=" + repr(train_dataset[1][0].size()))
    42. print("train_dataset.class_to_idx=" + repr(train_dataset.class_to_idx))
    43. # 创建训练集的可迭代对象，一个batch_size地读取数据,shuffle设为True表示随机打乱顺序读取
    44. train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    45.  
    46. # 测试集
    47. val_dataset = datasets.ImageFolder(root=base_dir_val, transform=pipeline)
    48. print(val_dataset)
    49. print("val_dataset=" + repr(val_dataset[1][0].size()))
    50. print("val_dataset.class_to_idx=" + repr(val_dataset.class_to_idx))
    51. # 创建测试集的可迭代对象，一个batch_size地读取数据
    52. val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=True)
    53.  
    54.  
    55. # 获得一批测试集的数据
    56. images, labels = next(iter(val_loader))
    57. print(images.shape)
    58. print(labels.shape)
    59.  
    60.  
    61. # 损失函数,交叉熵损失函数
    62. criterion = nn.CrossEntropyLoss()
    63.  
    64. # 使用预训练模型
    65. resnet_model = torchvision.models.resnet50(pretrained=True)
    66. num_ftrs = resnet_model.fc.in_features
    67. resnet_model.fc = nn.Linear(num_ftrs, 10)
    68. resnet_model.to(DEVICE)
    69. # 选择简单暴力的Adam优化器，学习率调低
    70. optimizer = optim.Adam(resnet_model.parameters(), lr=modellr)
    71. #optimizer = optim.SGD(net.parameters(), lr = 0.01)
    72.  
    73. train_loss_list = []
    74. train_accuracy_list = []
    75. test_loss_list = []
    76. test_accuracy_list = []
    77. train_iteration_list = []
    78. test_iteration_list = []
    79.  
    80.  
    81. best_val_acc = 0
    82.  
    83.  
    84. # 定义训练方法
    85. def train(model, device, train_loader, optimizer, epoch):
    86.     iteration = 0
    87.     train_correct = 0.0
    88.     model.train()
    89.     sum_loss = 0.0
    90.     total_num = len(train_loader.dataset)
    91.     print(total_num, len(train_loader))
    92.     for batch_idx, (data, target) in enumerate(train_loader):
    93.         # 获取数据与标签
    94.         data, target = Variable(data).to(device), Variable(target).to(device)
    95.         
    96.         # 梯度清零
    97.         optimizer.zero_grad()
    98.         
    99.         # 计算损失
    100.         output = model(data)
    101.         loss = criterion(output, target)
    102.         
    103.         #反向传播
    104.         loss.backward()
    105.         
    106.         #更新参数
    107.         optimizer.step()
    108.         
    109.         print_loss = loss.data.item()
    110.         sum_loss += print_loss
    111.         _, train_predict = torch.max(output.data, 1)
    112.  
    113.         if torch.cuda.is_available():
    114.             train_correct += (train_predict.cuda() == target.cuda()).sum()
    115.         else:
    116.             train_correct += (train_predict == target).sum()
    117.         accuracy = (train_correct / total_num) * 100
    118.         print("Epoch: %d , Batch: %3d , Loss : %.8f,train_correct:%d , train_total:%d , accuracy:%.6f" % (
    119.             epoch + 1, batch_idx + 1, loss.item(), train_correct, total_num, accuracy))
    120.         # 存在集合画图
    121.         if (epoch + 1) == EPOCHS:  # 只画出最后一个epoch时候的准确度变化曲线
    122.             iteration += 1
    123.             train_loss_list.append(loss.item())
    124.             train_iteration_list.append(iteration)
    125.             train_accuracy_list.append(accuracy)
    126.  
    127.  
    128. # 定义验证方法
    129. def val(model, device, val_loader, epoch):
    130.     print("=====================预测开始=================================")
    131.     iteration = 0
    132.     model.eval()
    133.     test_loss = 0.0
    134.     correct = 0.0
    135.     total_num = len(val_loader.dataset)
    136.     print(total_num, len(val_loader))
    137.     with torch.no_grad():
    138.         for data, target in val_loader:
    139.             data, target = Variable(data).to(device), Variable(target).to(device)
    140.             output = model(data)
    141.             loss = criterion(output, target)
    142.             _, pred = torch.max(output.data, 1)
    143.             if torch.cuda.is_available():
    144.                 correct += torch.sum(pred.cuda() == target.cuda())
    145.             else:
    146.                 correct += torch.sum(pred == target)
    147.             print_loss = loss.data.item()
    148.             test_loss += print_loss
    149.         acc = correct / total_num * 100
    150.         avg_loss = test_loss / len(val_loader)
    151.         """
    152.             因为调用这个方法的时候就是每次结束训练一次之后调用
    153.         """
    154.         # iteration += 1
    155.         # 存入集合准备画图
    156.         test_loss_list.append(avg_loss)
    157.         test_accuracy_list.append(acc)
    158.         test_iteration_list.append(epoch)
    159.         print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.6f}%)\n'.format(
    160.             avg_loss, correct, len(val_loader.dataset), acc))
    161.             
    162.         global best_val_acc
    163.         if acc > best_val_acc:
    164.             best_val_acc = acc
    165.             print("Best Accuracy:{:.6f}%".format(best_val_acc))
    166.             torch.save(resnet_model.state_dict(), 'best-{:.6f}.model.pth'.format(best_val_acc))  # 保存模型
    167.  
    168.  
    169. # 训练
    170. for epoch in range(EPOCHS):
    171.     train(resnet_model, DEVICE, train_loader, optimizer, epoch)
    172.     val(resnet_model, DEVICE, val_loader, epoch)
    173.     #torch.save(resnet_model, 'model.pth')  # 保存模型
    174.  
    175. # 可视化测试机的loss和accuracy
    176. plt.figure(1)
    177. plt.plot(test_iteration_list, test_loss_list)
    178. plt.title("ResNet50 test loss")
    179. plt.ylabel("loss")
    180. plt.xlabel("Number of test iteration")
    181. plt.show()
    182.  
    183. plt.figure(2)
    184. plt.plot(test_iteration_list, test_accuracy_list)
    185. plt.title("ResNet50 test accuracy")
    186. plt.xlabel("Number of test iteration")
    187. plt.ylabel("accuracy")
    188. plt.show()
    189.  
    190. # 可视化训练集loss和accuracy
    191. plt.figure(3)
    192. plt.plot(train_iteration_list, train_loss_list)
    193. plt.title("ResNet50 train loss")
    194. plt.xlabel("Number of train iteration")
    195. plt.ylabel("accuracy")
    196. plt.show()
    197.  
    198. plt.figure(4)
    199. plt.plot(train_iteration_list, train_accuracy_list)
    200. plt.title("ResNet50 train accuracy")
    201. plt.xlabel("Number of train iteration")
    202. plt.ylabel("accuracy")
    203. plt.show()

    代码需要注意的是数据集路径，用的是绝对路径，自行修改。

    代码训练的epoch是15，等待一段时间吧！

    五、测试模型

    测试模型脚本predict.py

    1. import os
    2. from PIL import Image
    3. import cv2
    4. import torch
    5. import torch.nn as nn
    6. from torch.autograd.variable import Variable
    7. import torchvision
    8. from torchvision import transforms
    9.  
    10. # 0-SUV, 1-BUS, 2-family sedan, 3-fire engine, 4-heavy truck, 
    11. # 5-jeep, 6-mini bus, 7-racing car, 8-taxi, 9-truck
    12.  
    13. def predict_single_image():
    14.     
    15.     MODEL_SAVE_FILE = 'best-82.000000.model.pth'
    16.     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    17.     
    18.     model = torchvision.models.resnet50()
    19.     num_ftrs = model.fc.in_features
    20.     model.fc = nn.Linear(num_ftrs, 10)
    21.     model.to(device)
    22.     
    23.     model.load_state_dict(torch.load(MODEL_SAVE_FILE,map_location='cpu'))
    24.     
    25.     
    26.     model = torch.nn.DataParallel(model,device_ids=[0])
    27.     model.eval()
    28.     
    29.     img = cv2.imread("test.jpg")
    30.     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    31.     image = Image.fromarray(img)
    32.     
    33.     pipeline = transforms.Compose([
    34.     # 分辨率重置为256
    35.     transforms.Resize(256),
    36.     # 对加载的图像作归一化处理， 并裁剪为[224x224x3]大小的图像(因为这图片像素不一致直接统一)
    37.     transforms.CenterCrop(224),
    38.     # 将图片转成tensor
    39.     transforms.ToTensor(),
    40.     # 正则化，模型出现过拟合现象时，降低模型复杂度
    41.     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    42.     ])
    43.  
    44.     image = pipeline(image)
    45.     image = image.unsqueeze(0)
    46.     print(image.shape)
    47.  
    48.     input_var = Variable(image).float().to(device)
    49.     output = model(input_var)
    50.     print("output:", output)
    51.     print("output.shape:", output.shape)
    52.     
    53.     soft_output = torch.softmax(output, dim=-1)
    54.     print("soft_output:", soft_output)
    55.     
    56.     percent, predicted = torch.max(soft_output.data, 1)
    57.     print("percent:", percent)
    58.     print("predicted:", predicted)
    59.     
    60.     '''
    61.     USE_GPU = torch.cuda.is_available()
    62.     if USE_GPU:
    63.         inputs = inputs.cuda()
    64.     if not os.path.exists(MODEL_SAVE_FILE):
    65.         print('can not find model save file.')
    66.         exit()
    67.     else:
    68.         if USE_GPU:
    69.             model.load_state_dict(torch.load(MODEL_SAVE_FILE))
    70.         else:
    71.             model.load_state_dict(torch.load(MODEL_SAVE_FILE, map_location=lambda storage, loc: storage))
    72.         outputs = model(inputs)
    73.         _, prediction_tensor = torch.max(outputs.data, 1)
    74.         if USE_GPU:
    75.             prediction = prediction_tensor.cpu().numpy()[0][0]
    76.             print('predict: ', prediction)
    77.             print('this is {}'.format(classes_name[prediction]))
    78.         else:
    79.             prediction = prediction_tensor.numpy()[0][0]
    80.             print('predict: ', prediction)
    81.             print('this is {}'.format(classes_name[prediction]))
    82.     '''
    83.  
    84.  
    85. predict_single_image()

    运行

    python predict.py

    六、模型转换

    1、转成onnx模型

    pth_to_onnx.py

    1. import torch
    2. import torch.nn as nn
    3. import torchvision
    4. from torch.autograd.variable import Variable
    5.  
    6.  
    7. MODEL_SAVE_FILE = 'best-82.000000.model.pth'
    8. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    9.  
    10. model = torchvision.models.resnet50()
    11. num_ftrs = model.fc.in_features
    12. model.fc = nn.Linear(num_ftrs, 10)
    13. model.to(device)
    14.  
    15. model.load_state_dict(torch.load(MODEL_SAVE_FILE,map_location='cpu'))
    16.  
    17. batch_size = 1  #批处理大小
    18.  
    19. # #set the model to inference mode
    20. model.eval()
    21.  
    22. d_input = Variable(torch.randn(1, 3, 224, 224))
    23. export_onnx_file = "10class_ResNet50.onnx"		# 目的ONNX文件名
    24. torch.onnx.export(model, d_input, export_onnx_file, opset_version=12,verbose=True)

    这里需要注意的 是opset_version算子，rk3568用12

    python pth_to_onnx.py

    onnx模型是我需要的，打算部署到rk3568，需要把onnx模型转成rknn模型，后续测试

    2、转成pt模型

    pth_to_pt.py

    1. import torch
    2. import torch.nn as nn
    3. import torchvision
    4.  
    5.  
    6. MODEL_SAVE_FILE = 'best-82.000000.model.pth'
    7. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    8.  
    9. model = torchvision.models.resnet50()
    10. num_ftrs = model.fc.in_features
    11. model.fc = nn.Linear(num_ftrs, 10)
    12. model.to(device)
    13.  
    14. model.load_state_dict(torch.load(MODEL_SAVE_FILE,map_location='cpu'))
    15.  
    16. model.eval()
    17.  
    18. example = torch.rand(1,3,224,224).to(device)
    19. traced_script_module = torch.jit.trace(model, example)
    20. traced_script_module.save('./10class_ResNet50.pt')

    运行转换：

    python pth_to_pt.py

    如有侵权，或需要完整代码，请及时联系博主。

  • 相关阅读:
    【Qt】Unicode编码作用 ，以及在Qt中的理解
    java毕业设计民宿管理平台mybatis+源码+调试部署+系统+数据库+lw
    第30章 市场营销
    R语言实战应用案例：论文篇（二）-特殊盒须图绘制
    spring框架漏洞整理(Spring Data漏洞)
    Redis整理
    SpringBoot集成Mybatis-Plus
    TCP协议
    日更【系统架构设计师知识总结2】指令系统（结合真题）
    使用Python进行钻石价格分析
  • 原文地址：https://blog.csdn.net/weixin_38807927/article/details/133915602