【PyTorch笔记】pytorch入门教程14 网络模型的使用、修改、保存和读取

最新推荐文章于 2024-09-13 23:12:01 发布

鱼鱼9901

最新推荐文章于 2024-09-13 23:12:01 发布

阅读量119

点赞数

分类专栏： pytorch 文章标签： pytorch 笔记深度学习

本文链接：https://blog.csdn.net/weixin_72100405/article/details/130667181

版权

pytorch 专栏收录该内容

17 篇文章 0 订阅

订阅专栏

该文展示了如何在PyTorch中使用预训练的VGG16模型，并添加自定义线性层。同时，讨论了两种模型保存与加载的方法，包括保存整个模型和仅保存状态字典。

摘要由CSDN通过智能技术生成

import torch
import torchvision.models
from torch import nn

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True)

print(vgg16_true)

train_data = torchvision.datasets.CIFAR10('./dataset',train = True,transform=torchvision.transforms.ToTensor())
vgg16_true.add_module('add_linera1',nn.Linear(1000,10))
# vgg16_true.classifier.add_module('add_module1',nn.Linear(1000,10))
# 第一个add是在整个模型所有Module后面加一个线性层，输入1000通道输出10通道，第二个是在原来名为classifier的module后面加一个线性层
print(vgg16_true)

print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096,10)
#把classifiermodule的第六个改成这个线性层
print(vgg16_false)

# 保存与加载
vgg16 = vgg16_false
torch.save(vgg16,'vgg16_method1.pth')
#方法一保存模型结构和参数
torch.save(vgg16.state_dict(),'vgg16_method2.pth')
#方法二只把参数保存成字典形式


model = torch.load('vgg16_method1.pth')#（）里是路径
#方法一在load的时候要from 原py文件名 import *；或者把原来创定的Tudui复制过来

vgg16 = torchvision.models.vgg16(pretrained = False)
vgg16.load_state_dic(torch.load('vgg16_method2.pth'))
#方法二先加载原网络模型，再加载自己的参数
# model = torch.load('vgg16_method2.pth') 这是方法二的只加载参数