代码实战（小土堆）——网络模型的保存与读取

方式一保存

import torchvision
import torch
from torch import nn

vgg16 = torchvision.models.vgg16(weights=None)
# 保存方式1（模型的结构和参数都被保存）
torch.save(vgg16, "vgg16_method1.pth")

方式一读取

import torch
import torchvision

# 方式1 ——> 保存方式1，加载模型
model = torch.load("vgg16_method1.pth")
print(model)

 

方式二保存

import torchvision
import torch
from torch import nn

vgg16 = torchvision.models.vgg16(weights=None)
# 保存方式2（保存成字典形式，只保存参数）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

方式二读取

import torch
import torchvision

# 方式2 ——> 保存方式2，加载模型
vgg16 = torchvision.models.vgg16(weights=None)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))

# 打印出来的只有模型参数
# model2 = torch.load("vgg16_method2.pth")

print(vgg16)

陷阱

import torchvision
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, "tudui_method1.pth")

读取

import torch
import torchvision

# 陷阱1(方式1)
model = torch.load("tudui_method1.pth")
print(model)

 报错了。。。

 用方式一保存自己的网络模型，加载的时候需要复制过去。。。

import torch
import torchvision
from torch import nn

# 陷阱1(方式1)
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
    
model = torch.load("tudui_method1.pth")
print(model)

但是在真实的项目当中不会把模型复制来复制去，会定义在一个单独的文件夹里面，通过import引入过来

import torch
import torchvision
from torch import nn
from  model_save import *

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