[土堆]深度学习快速入门教程笔记——网络模型保存与读取——015

1. 网络模型保存

---

1.1保存方式一

import torchvision
import torch
vgg16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg16,"./model/vgg16_method1.pth") # 保存方式一：模型结构 + 模型参数      
print(vgg16)

1.2保存方式二

import torchvision
import torch
vgg16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg16.state_dict(),"./model/vgg16_method2.pth")
# 保存方式二：模型参数（官方推荐）,不再保存网络模型结构  
print(vgg16)

---

2.网络模型加载

---

2.1加载方式一（保存方式一对应的加载模型）

import torch
model = torch.load("./model/vgg16_method1.pth") 
# 保存方式一对应的加载模型    
print(model)

2.2加载方式二（保存方式二对应的加载模型）

import torch
import torchvision
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_model2.pth"))  
# 输出完整的模型结构，与第一种方式输出的模型结构相同
print(vgg16)

---

3.网络陷阱

---

现有自定义网络模型并保存：

import torch
from torch import nn

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui,self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)
        
    def forward(self,x):
        x = self.conv1(x)
        return x

tudui = Tudui()
torch.save(tudui, "./model/tudui_method1.pth")

再运行下面的代码，即下面为第1个代码块运行，无法直接导入网络模型。

import torch
model = torch.load("./model/tudui_method1.pth")  # 无法直接加载方式一保存的网络结构    
print(model)

解决方法1：

import torch
from torch import nn

# 确保网络模型是我们想要的网络模型,要在加载前还写明网络模型
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui,self).__init__()
        self.conv1 = nn.Conv2d(3,64,kernel_size=3)
        
    def forward(self,x):
        x = self.conv1(x)
        return x
    
#tudui = Tudui # 不需要写这一步，不需要创建网络模型    
model = torch.load("./model/tudui_method1.pth")  # 无法直接加载方式一保存的网络结构    
print(model)

解决方法2：

import torch
import model_save import * # 它就相当于把 model_save.py 里的网络模型定义写到这里了
    
#tudui = Tudui # 不需要写这一步，不需要创建网络模型    

model = torch.load("tudui_method1.pth")
print(model)