一、案例

1.保存方式1

保存模型结构+模型参数

2.保存方式1的陷阱

自己创建的模型，在加载的时候，需要import

3.保存方式2

保存模型参数（官方推荐）

比较1和2两种方式保存的内容的大小
terminal执行 ls -all

4.完整代码

• 模型保存

import torch
import torchvision
from torch import nn

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

#保存方式2：模型参数（官方推荐）
torch.save(vgg16.state_dict(),"vgg16_method2.pth")

#保存方式1的陷阱：自己创建的模型，在调用的时候需要import,不然会报错
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conva=nn.Conv2d(3,64,kernel_size=3)

    def forward(self,x):
        x=self.conv1(x)
        return x

tudui=Tudui()
torch.save(tudui,"tudui_method.pth")

• 模型加载

import torch

# 方式1-》加载模型
import torchvision

model=torch.load("vgg16_method1.pth")
# print(model)

# 方式2-》通过加载模型参数加载模型
vgg16=torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# print(vgg16)

# 方式1的陷阱
from test26_1_model_save import * #引入对应模型的结构
model1=torch.load("tudui_method.pth")
print(model1)

标签：26,读取,pth,vgg16,模型,torch,保存,import
来源： https://www.cnblogs.com/yuyingblogs/p/16308652.html