【Pytorch】12.网络模型的加载、修改与保存

Elephant_King

于 2024-05-17 14:17:54 发布

阅读量505

点赞数 6

CC 4.0 BY-SA版权

分类专栏： Pytorch 文章标签： pytorch 网络人工智能

本文链接：https://blog.csdn.net/Elephant_King/article/details/139002934

已有网络模型的加载与保存

对于我们要保存的网络模型，有两种保存方法

保存网络模型的结构与参数
仅保留网络模型的参数

import torch
import torchvision
from torch import nn

# 训练好的vgg网络
vgg_train = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.IMAGENET1K_V1)
# vgg_test = torchvision.models.vgg16(pretrained=True)


# 未训练的vgg网络，只有网络结构
vgg_not_train = torchvision.models.vgg16(weights=None)
# vgg_not_train = torchvision.models.vgg16(pretrained=True)

# 保存神经网络,保存所有信息(结构+参数)
torch.save(vgg_train, 'vgg16_method1.pth')
# 保存神经网络参数
torch.save(vgg_train.state_dict(), 'vgg16_method2.pth')

print(vgg_train

最低0.47元/天解锁文章