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

pytorch实现resnet50（训练+测试+模型转换）_pytorch resnet50预训练模型

作者：花生_TL007 | 2024-04-06 10:17:13

赞

踩

pytorch resnet50预训练模型

本章使用pytorch训练resnet50，使用cifar数据集。

数据集：

代码工程：

1.train.py

  1.  
  2. import torch
  3. from torch import nn, optim
  4. import torchvision.transforms as transforms
  5. from torchvision import datasets
  6. from torch.utils.data import DataLoader
  7. from resnet50 import ResNet50
  8.  
  9.  
  10. #  用CIFAR-10 数据集进行实验
  11.  
  12. def main():
  13.     batchsz = 2
  14.  
  15.     cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
  16.         transforms.Resize((224, 224)),
  17.         transforms.ToTensor(),
  18.         transforms.Normalize(mean=[0.485, 0.456, 0.406],
  19.                              std=[0.229, 0.224, 0.225])
  20.     ]), download=True)
  21.     cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
  22.     
  23.     cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
  24.         transforms.Resize((224, 224)),
  25.         transforms.ToTensor(),
  26.         transforms.Normalize(mean=[0.485, 0.456, 0.406],
  27.                              std=[0.229, 0.224, 0.225])
  28.     ]), download=True)
  29.     cifar_test = DataLoader(cifar_test, batch_size=1, shuffle=True)
  30.  
  31.     x, label = iter(cifar_train).next()
  32.     print('x:', x.shape, 'label:', label.shape)
  33.  
  34.     device = torch.device('cpu')
  35.     model = ResNet50().to(device)
  36.     print(*list(model.children())[-3:-2])
  37.  
  38.  
  39.     criteon = nn.CrossEntropyLoss().to(device)
  40.     optimizer = optim.Adam(model.parameters(), lr=1e-3)
  41.     # print(model)
  42.  
  43.     print(iter(cifar_test).next()[0].shape)
  44.  
  45.     for epoch in range(1):
  46.         model.train()
  47.  
  48.         for batchidx, (x, label) in enumerate(cifar_train):
  49.             if batchidx<=2:
  50.                 x, label = x.to(device), label.to(device)
  51.  
  52.                 logits = model(x)
  53.                 loss = criteon(logits, label)
  54.  
  55.                 # backprop
  56.                 optimizer.zero_grad()
  57.                 loss.backward()
  58.                 optimizer.step()
  59.                 print("epoch:",epoch, "index:",batchidx)
  60.             else:
  61.                 continue
  62.         print(epoch, 'loss:', loss.item())
  63.  
  64.     #     # # PATH="model/test.pth"
  65.         torch.save(model, "model2/test.pth")
  66.         torch.save(model.state_dict(),"model2/test2.pth")
  67.  
  68.         model.eval()
  69.         with torch.no_grad():
  70.             # test
  71.             total_correct = 0
  72.             total_num = 0
  73.             for idx, (x, label) in enumerate(cifar_test):
  74.                 if idx<=5:
  75.                     x, label = x.to(device), label.to(device)
  76.  
  77.                     logits = model(x)
  78.                     pred = logits.argmax(dim=1)
  79.                     
  80.                     correct = torch.eq(pred, label).float().sum().item()
  81.                     total_correct += correct
  82.                     total_num += x.size(0)
  83.                     # print(pred)
  84.  
  85.             acc = total_correct / total_num
  86.             print(epoch, 'test acc:', acc)
  87.  
  88.  
  89. if __name__ == '__main__':
  90.     main()
  91.  
  92.  
  93.  
  94.  
  95. # # 保存整个网络
  96. # torch.save(net, PATH) 
  97. # # 保存网络中的参数, 速度快，占空间少
  98. # torch.save(net.state_dict(),PATH)
  99. # #--------------------------------------------------
  100. # #针对上面一般的保存方法，加载的方法分别是：
  101. # model_dict=torch.load(PATH)
  102. # model_dict=model.load_state_dict(torch.load(PATH))

2.test_pth.py

  1. from resnet50 import ResNet50
  2. import torch
  3. from PIL import Image
  4. from torchvision import transforms
  5. import cv2
  6. import numpy as np
  7.  
  8. def prediect(img_path):
  9.     device = torch.device('cpu')
  10.     model=torch.load("model2/test.pth")
  11.     model=model.to(device)
  12.     
  13.     # model = ResNet50()
  14.     # weight=torch.load("model/test2.pth")
  15.     # model.load_state_dict(weight)
  16.     # model=model.to(device)
  17.  
  18.     img=cv2.imread(img_path)
  19.     img=cv2.resize(img, (224, 224))
  20.     img=np.reshape(img,(1,224,224,3))
  21.     img=img.transpose(0,3,1,2).copy()
  22.     print(img.shape)
  23.     img_ = torch.Tensor(img)
  24.     
  25.     torch.no_grad()
  26.     outputs = model(img_)
  27.     _, predicted = torch.max(outputs, 1)
  28.     print('pred :',outputs, predicted)
  29.  
  30. if __name__ == '__main__':
  31.     img_path="img/dog2.jpg"
  32.     prediect(img_path)
  33.

3.resnet50.py

  1. import torch
  2. import torch.nn as nn
  3. from torch.nn import functional as F
  4.  
  5.  
  6. class ResNet50BasicBlock(nn.Module):
  7.     def __init__(self, in_channel, outs, kernerl_size, stride, padding):
  8.         super(ResNet50BasicBlock, self).__init__()
  9.         self.conv1 = nn.Conv2d(in_channel, outs[0], kernel_size=kernerl_size[0], stride=stride[0], padding=padding[0])
  10.         self.bn1 = nn.BatchNorm2d(outs[0])
  11.         self.conv2 = nn.Conv2d(outs[0], outs[1], kernel_size=kernerl_size[1], stride=stride[0], padding=padding[1])
  12.         self.bn2 = nn.BatchNorm2d(outs[1])
  13.         self.conv3 = nn.Conv2d(outs[1], outs[2], kernel_size=kernerl_size[2], stride=stride[0], padding=padding[2])
  14.         self.bn3 = nn.BatchNorm2d(outs[2])
  15.  
  16.     def forward(self, x):
  17.         out = self.conv1(x)
  18.         out = F.relu(self.bn1(out))
  19.  
  20.         out = self.conv2(out)
  21.         out = F.relu(self.bn2(out))
  22.  
  23.         out = self.conv3(out)
  24.         out = self.bn3(out)
  25.  
  26.         return F.relu(out + x)
  27.  
  28.  
  29. class ResNet50DownBlock(nn.Module):
  30.     def __init__(self, in_channel, outs, kernel_size, stride, padding):
  31.         super(ResNet50DownBlock, self).__init__()
  32.         # out1, out2, out3 = outs
  33.         # print(outs)
  34.         self.conv1 = nn.Conv2d(in_channel, outs[0], kernel_size=kernel_size[0], stride=stride[0], padding=padding[0])
  35.         self.bn1 = nn.BatchNorm2d(outs[0])
  36.         self.conv2 = nn.Conv2d(outs[0], outs[1], kernel_size=kernel_size[1], stride=stride[1], padding=padding[1])
  37.         self.bn2 = nn.BatchNorm2d(outs[1])
  38.         self.conv3 = nn.Conv2d(outs[1], outs[2], kernel_size=kernel_size[2], stride=stride[2], padding=padding[2])
  39.         self.bn3 = nn.BatchNorm2d(outs[2])
  40.  
  41.         self.extra = nn.Sequential(
  42.             nn.Conv2d(in_channel, outs[2], kernel_size=1, stride=stride[3], padding=0),
  43.             nn.BatchNorm2d(outs[2])
  44.         )
  45.  
  46.     def forward(self, x):
  47.         x_shortcut = self.extra(x)
  48.         out = self.conv1(x)
  49.         out = self.bn1(out)
  50.         out = F.relu(out)
  51.  
  52.         out = self.conv2(out)
  53.         out = self.bn2(out)
  54.         out = F.relu(out)
  55.  
  56.         out = self.conv3(out)
  57.         out = self.bn3(out)
  58.         return F.relu(x_shortcut + out)
  59.  
  60.  
  61. class ResNet50(nn.Module):
  62.     def __init__(self):
  63.         super(ResNet50, self).__init__()
  64.         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
  65.         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
  66.  
  67.         self.layer1 = nn.Sequential(
  68.             ResNet50DownBlock(64, outs=[64, 64, 256], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0]),
  69.             ResNet50BasicBlock(256, outs=[64, 64, 256], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0]),
  70.             ResNet50BasicBlock(256, outs=[64, 64, 256], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0]),
  71.         )
  72.  
  73.         self.layer2 = nn.Sequential(
  74.             ResNet50DownBlock(256, outs=[128, 128, 512], kernel_size=[1, 3, 1], stride=[1, 2, 1, 2], padding=[0, 1, 0]),
  75.             ResNet50BasicBlock(512, outs=[128, 128, 512], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0]),
  76.             ResNet50BasicBlock(512, outs=[128, 128, 512], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0]),
  77.             ResNet50DownBlock(512, outs=[128, 128, 512], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1], padding=[0, 1, 0])
  78.         )
  79.  
  80.         self.layer3 = nn.Sequential(
  81.             ResNet50DownBlock(512, outs=[256, 256, 1024], kernel_size=[1, 3, 1], stride=[1, 2, 1, 2], padding=[0, 1, 0]),
  82.             ResNet50BasicBlock(1024, outs=[256, 256, 1024], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1],
  83.                                padding=[0, 1, 0]),
  84.             ResNet50BasicBlock(1024, outs=[256, 256, 1024], kernerl_size=[1, 3, 1], stride=[1, 1, 1, 1],
  85.                                padding=[0, 1, 0]),
  86.             ResNet50DownBlock(1024, outs=[256, 256, 1024], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1],
  87.                               padding=[0, 1, 0]),
  88.             ResNet50DownBlock(1024, outs=[256, 256, 1024], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1],
  89.                               padding=[0, 1, 0]),
  90.             ResNet50DownBlock(1024, outs=[256, 256, 1024], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1],
  91.                               padding=[0, 1, 0])
  92.         )
  93.  
  94.         self.layer4 = nn.Sequential(
  95.             ResNet50DownBlock(1024, outs=[512, 512, 2048], kernel_size=[1, 3, 1], stride=[1, 2, 1, 2],
  96.                               padding=[0, 1, 0]),
  97.             ResNet50DownBlock(2048, outs=[512, 512, 2048], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1],
  98.                               padding=[0, 1, 0]),
  99.             ResNet50DownBlock(2048, outs=[512, 512, 2048], kernel_size=[1, 3, 1], stride=[1, 1, 1, 1],
  100.                               padding=[0, 1, 0])
  101.         )
  102.  
  103.         self.avgpool = nn.AvgPool2d(kernel_size = 7,stride=1,ceil_mode=False)
  104.         # self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
  105.  
  106.         self.fc = nn.Linear(2048, 10)
  107.         # 使用卷积代替全连接
  108.         self.conv11=nn.Conv2d(2048, 10, kernel_size=1, stride=1, padding=0)
  109.  
  110.     def forward(self, x):
  111.         out = self.conv1(x)
  112.         out = self.maxpool(out)
  113.         out = self.layer1(out)
  114.         out = self.layer2(out)
  115.         out = self.layer3(out)
  116.         out = self.layer4(out)
  117.         out = self.avgpool(out)
  118.         out=self.conv11(out)
  119.         out = out.reshape(x.shape[0], -1)
  120.         # out = self.fc(out)
  121.         return out
  122.  
  123.  
  124. if __name__ == '__main__':
  125.     x = torch.randn(1, 3, 224, 224)
  126.     net = ResNet50()
  127.     out = net(x)
  128.     print('out.shape: ', out.shape)
  129.     print(out)
  130.

4.pth2onnx.py

  1. import torch
  2. from torchsummary import summary
  3.  
  4. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  5.  
  6. model = torch.load("model2/test.pth") # pytorch模型加载
  7. model.eval()
  8. for name in model.state_dict():
  9. 	print(name)
  10. summary(model, (3, 224, 224))
  11.  
  12. input_shape=list(map(int, "1,3,224,224".split(",")))
  13. x = torch.randn(input_shape)   # 生成张量
  14. x = x.to(device)
  15.  
  16. export_onnx_file = "model2/test.onnx"		# 目的ONNX文件名
  17. torch.onnx.export(model, x, export_onnx_file, verbose=True)
  18. # torch.onnx.export(model, x, export_onnx_file, verbose=True, export_params=True, do_constant_folding=True, opset_version=11)
  19.  
  20.  
  21. # input_names=['boxes']
  22. # output_names=['layer1.1.conv1.bias']
  23. # torch.onnx.export(model, x, export_onnx_file,
  24. # 					export_params=True,
  25. # 					do_constant_folding=True,
  26. # 					input_names=input_names, 
  27. # 					output_names=output_names
  28. # 					)

5.test_onnx_v1.py

  1. import cv2
  2. import numpy as np
  3. import onnxruntime as rt
  4.  
  5. def image_process(image_path):
  6.     mean = np.array([[[0.485, 0.456, 0.406]]])      # 训练的时候用来mean和std
  7.     std = np.array([[[0.229, 0.224, 0.225]]])
  8.  
  9.     img = cv2.imread(image_path)
  10.     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  11.     img = cv2.resize(img, (224, 224))                 # (96, 96, 3)
  12.  
  13.     image = img.astype(np.float32)/255.0
  14.     image = (image - mean)/ std
  15.  
  16.     image = image.transpose((2, 0, 1))              # (3, 96, 96)
  17.     image = image[np.newaxis,:,:,:]                 # (1, 3, 96, 96)
  18.  
  19.     image = np.array(image, dtype=np.float32)
  20.     
  21.     return image
  22.  
  23. def onnx_runtime():
  24.     imgdata = image_process('img/test.jpg')
  25.     
  26.     sess = rt.InferenceSession("model2/test.onnx")
  27.     input_name = sess.get_inputs()[0].name  
  28.     output_name = sess.get_outputs()[0].name
  29.  
  30.     pred_onnx = sess.run([output_name], {input_name: imgdata})
  31.  
  32.     print("outputs:",np.array(pred_onnx)[0].shape)
  33.  
  34. onnx_runtime()

6.test_onnx_v2.py

  1. import numpy as np
  2. import torch
  3. import onnx
  4. import onnxruntime
  5. import pickle
  6.  
  7. # 测试数据
  8. x = torch.randn(1,3,224,224, requires_grad=False)
  9. print(type(x),x.shape)
  10. # 使用 ONNX 的 API 检查 ONNX 模型
  11. onnx_model = onnx.load("model2/test.onnx")
  12. onnx.checker.check_model(onnx_model)
  13.  
  14. # onnx模型测试
  15. ort_session = onnxruntime.InferenceSession("model2/test.onnx")
  16. def to_numpy(tensor):
  17.     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
  18.  
  19.  #结果输出
  20. ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
  21. ort_outs = ort_session.run(None, ort_inputs)
  22. ort_out = ort_outs[0]
  23. print(x.shape, ort_out.shape)
  24.  
  25.  
  26. # torch模型测试
  27. # model=torch.load("test/person_reid.pth",map_location='cpu')
  28. # model.eval()
  29. # torch_out = model(x)
  30.  
  31. # 比较ONNX 和 PyTorch 的结果
  32. # np.testing.assert_allclose(to_numpy(torch_out), ort_out, rtol=1e-03, atol=1e-05)
  33. # print("模型没有太大差异!")

7.onnx2pb.py

  1. import onnx
  2. from onnx_tf.backend import prepare
  3.  
  4. def onnx2pb(onnx_input_path, pb_output_path):
  5.     onnx_model = onnx.load(onnx_input_path)  # load onnx model
  6.     tf_exp = prepare(onnx_model)  # prepare tf representation
  7.     tf_exp.export_graph(pb_output_path)  # export the model
  8.  
  9. if __name__ == "__main__":
  10.     # onnx_input_path = 'test/person_reid.onnx'
  11.     # pb_output_path = 'test/person_reid2.pb'
  12.     
  13.     onnx_input_path = 'model2/test.onnx'
  14.     pb_output_path = 'model2/test.pb'
  15.  
  16.     onnx2pb(onnx_input_path, pb_output_path)

8.test_pb.py  (onnx+pb)

  1. import tensorflow as tf
  2. from tensorflow.python.framework import graph_util
  3. from tensorflow.python import pywrap_tensorflow
  4. import cv2
  5. import numpy as np
  6. import torch
  7. import onnx
  8. import onnxruntime
  9. import pickle
  10.  
  11.  
  12.  
  13. def recognize(img, pb_file_path):
  14.     with tf.Graph().as_default():
  15.         output_graph_def = tf.GraphDef()
  16.  
  17.         with open(pb_file_path, "rb") as f:#主要步骤即为以下标出的几步,1、2步即为读取图
  18.             output_graph_def.ParseFromString(f.read())# 1.将模型文件解析为二进制放进graph_def对象
  19.             _ = tf.import_graph_def(output_graph_def, name="")# 2.import到当前图
  20.  
  21.         with tf.Session() as sess:
  22.             init = tf.global_variables_initializer()
  23.             sess.run(init)
  24.  
  25.             graph = tf.get_default_graph()# 3.获得当前图
  26.             
  27.             # # 4.get_tensor_by_name获取需要的节点
  28.             # x = graph.get_tensor_by_name("IteratorGetNext_1:0")
  29.             # y_out = graph.get_tensor_by_name("resnet_v1_50_1/predictions/Softmax:0")
  30.  
  31.             x = graph.get_tensor_by_name("data:0")
  32.             y_out = graph.get_tensor_by_name("reid_embedding:0")
  33.             
  34.             # img=np.random.normal(size=(1, 224, 224, 3))
  35.             # img=cv2.imread(jpg_path)
  36.             # img=cv2.resize(img, (128, 256))
  37.             # img=np.reshape(img,(1,128,256,3))
  38.             # img=img.transpose(0,3,1,2).copy()
  39.             # print(img.shape)
  40.             
  41.             #执行
  42.             output = sess.run(y_out, feed_dict={x:img})
  43.             pred=np.argmax(output, axis=1)
  44.             return output
  45.             # print("预测结果：", output.shape, output, "预测label:", pred)
  46.  
  47.  
  48. jpg_path="img/test.jpg"
  49. img=cv2.imread(jpg_path)
  50. img=cv2.resize(img, (128, 256))
  51. img=np.reshape(img,(1,128,256,3))
  52. img=img.transpose(0,3,1,2).copy()
  53. print(img.shape)
  54. x = torch.randn(1,3,256,128, requires_grad=False)
  55. img=x
  56.  
  57. # 测试pb
  58. a=recognize(img, "test/gg.pb")
  59. print(a.shape)
  60. # b=recognize(img, "test/person_reid2.pb")
  61. # np.testing.assert_allclose(a, b, rtol=1e-03, atol=1e-05)
  62. # print(a.shape,a[0][4],b[0][4])
  63.  
  64.  
  65.  
  66. # # # 测试数据
  67. # # x = torch.randn(1,3,256,128, requires_grad=False)
  68. # # # x=torch.from_numpy(img)
  69. # # # x.requires_grad=False
  70.  
  71. # # 使用 ONNX 的 API 检查 ONNX 模型
  72. # onnx_model = onnx.load("test/person_reid.onnx")
  73. # onnx.checker.check_model(onnx_model)
  74.  
  75. # # onnx模型测试
  76. # ort_session = onnxruntime.InferenceSession("test/person_reid.onnx")
  77. # def to_numpy(tensor):
  78. #     return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
  79.  
  80. #  #结果输出
  81. # ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
  82. # ort_outs = ort_session.run(None, ort_inputs)
  83. # ort_out = ort_outs[0]
  84. # print(ort_out.shape, ort_out[0][4])
  85. # np.testing.assert_allclose(a, ort_out, rtol=1e-03, atol=1e-05)

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

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

  

闽ICP备14008679号