保存方式一 (保存模型结构和模型参数）
model_save.py
python
import torch
import torchvision
from torch import nn
vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1,模型结构+模型参数
torch.save(vgg16, "vgg16_method11.pth")model_load.py
python
import torch
from model_save import *
import torchvision
from torch import nn
model = torch.load("vgg16_method11.pth",weights_only=False)
print(model)
保存方式二（官方推荐方式）
model_save.py
state_dict: 将模型参数保存为Python字典格式
仅保存模型参数
加载时需要先重建模型结构再加载参数
model_save.py
python
import torch
import torchvision
from torch import nn
vgg16 = torchvision.models.vgg16(pretrained=False)
# # 保存方式2，模型参数（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")model_load.py
python
import torch
from model_save import *
import torchvision
from torch import nn
# 方式2，加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
model = torch.load("vgg16_method2.pth")
print(model) #打印出来是字典
print(vgg16) # 打印出来是模型
模型加载陷阱
问题现象: 使用方式一保存自定义模型后，加载时报错"Can't get attribute"
原因分析: Python的pickle机制需要能访问到原始类定义
解决方案一:
在加载文件中重新定义模型类
解决方案二:
将模型类定义放在单独模块中导入
