P26网络模型的保存与读取

/home/hcq/python/pytorch/module/01model_save

保存save

  2. import torch
  3. import torchvision
  4. from torch import nn
  6. vgg16 = torchvision.models.vgg16(pretrained=False)
  8. # 保存方式1: 模型结构+模型参数
  9. torch.save(vgg16, "../../pth/vgg16_method1.pth")
  11. # 保存方式2：（官方推荐）只保存参数，格式字典dict形式
  12. torch.save(vgg16.state_dict(), "../../pth/vgg16_method2.pth")
  15. # 陷阱
  16. class Mymodule(nn.Module):
  17.     def __init__(self):
  18.         super(Mymodule, self).__init__()
  19.         self.conv1 = nn.Conv2d(3,64,kernel_size=3)
  21.     def forward(self,x):
  22.         output = self.conv1(x)
  23.         return output
  25. mymodule = Mymodule()
  26. torch.save(mymodule, "../../pth/mymodule_method2.pth")
  27. print("导出完成")

加载load

  3. import torch
  5. # 方式1
  6. import torchvision.models
  7. from torch import nn
  9. model = torch.load("../../pth/vgg16_method1.pth")
  10. print(model)
  12. # VGG(
  13. #   (features): Sequential(
  14. #     (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  15. #     (1): ReLU(inplace=True)
  17. # 方式2: 输出是个字典
  18. model = torch.load("../../pth/vgg16_method2.pth")
  19. print(model)
  20. # OrderedDict([('features.0.weight', tensor([[[[ 0.0166, -0.0475,  0.0280],
  21. #           [-0.0281, -0.0843,  0.0382],
  22. #           [-0.0265, -0.1130, -0.0962]],
  24. ### 加载参数
  25. vgg16 = torchvision.models.vgg16(pretrained=False)
  26. vgg16.load_state_dict(torch.load("../../pth/vgg16_method2.pth"))
  27. print(vgg16)
  29. # 陷阱(解决方法：把class Mymodule 添加进来)
  30. ## 解决方法1
  31. class Mymodule(nn.Module):
  32.     def __init__(self):
  33.         super(Mymodule, self).__init__()
  34.         self.conv1 = nn.Conv2d(3,64,kernel_size=3)
  36.     def forward(self,x):
  37.         output = self.conv1(x)
  38.         return output
  39. ## 解决方法2
  40. # from 01model_save import *
  42. model = torch.load("../../pth/mymodule_method2.pth")
  43. print(model) # AttributeError: Can't get attribute 'Mymodule' on <module '__main__'