pytorch -- 加载和保存模型

最新推荐文章于 2024-04-30 15:11:08 发布

斑驳的岁月

最新推荐文章于 2024-04-30 15:11:08 发布

阅读量381

点赞数 11

分类专栏： pytorch 文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/weixin_39467231/article/details/136299003

版权

pytorch 专栏收录该内容

11 篇文章 1 订阅

订阅专栏

方法1： torch.load
保存网络模型结构和参数

import torch
model = torch.load('vgg16_method.pth')
torch.save(model,'vgg16_method_new.pth')

方法2：load_state_dict
保存网络模型中的参数为字典的形式（官方推荐）

import torch
import torchvision
# 加载没保存图结构的pth时，需要先初始化模型结构，即把架子搭好
vgg16_load = torchvision.models.vgg16(pretrained=False)
vgg16_load.load_state_dict(torch.load('vgg16_method2.pth'))
torch.save(vgg16_load.state_dict(),'vgg16_method2.pth')