【Pytorch】网络的保存与读取(6)

sdbhewfoqi

已于 2022-11-16 11:41:19 修改

阅读量682

点赞数

分类专栏： Tensorflow & Pytorch 文章标签： pytorch 网络深度学习

于 2021-06-29 10:46:48 首次发布

本文链接：https://blog.csdn.net/weixin_31866177/article/details/118313332

版权

Tensorflow & Pytorch 专栏收录该内容

41 篇文章 7 订阅

订阅专栏

2.2.3 从指定路径加载（而不是从默认路径加载）

2.3 方式一加载模型结构的陷阱

一、调用已存在的网络模型

使用vgg16,Models and pre-trained weights — Torchvision 0.14 documentation

False是初始化的参数，Ture是训练好的参数。

如下代码显示了，如何加载pytorch预训练的网络模型，并对网络进行修改。

vgg16原始输出结果是1000类，但是CIFAR10只有10类。需要对vgg16进行修改。

# -*- coding: utf-8 -*-

import torchvision

# train_data = torchvision.datasets.ImageNet("../data_image_net", split='train', download=True,
#                                            transform=torchvision.transforms.ToTensor())
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('../data', train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)

# 方法一：再增加一层，实现CIFAR10的10分类模型
vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)


# 方法二：修改vgg16网络的最后一层
# 将输出节点改为10
print(vgg16_false)
vgg16_false.classifier[6] = nn.Linear(4096, 10)
print(vgg16_false)

"""
也是ok！
vgg16_true.classifier[6] = nn.Linear(4096, 10)
print(vgg16_true)
"""

二、模型保存与加载

2.1 方式一——模型结构

2.1.1 保存

import torch
import torchvision
from torch import nn

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

2.1.2 加载

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

model = torch.load("vgg16_method1.pth")
print(model) # 打印模型结构

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

2.2 方式二——模型参数

2.2.1 保存

import torch
import torchvision
from torch import nn
# 保存方式2，模型参数（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth") # model参数，保存路径

2.2.2 加载

# 方式2，加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# model = torch.load("vgg16_method2.pth") # 这样直接加载，加载的是字典格式的参数数据，不是model结构
print(model) # 打印参数字典
print(vgg16) # 恢复成了model结构

2.2.3 从指定路径加载（而不是从默认路径加载）

    def _get_model_and_layer(self, model_name, layer):
        """ Internal method for getting layer from model
        :param model_name: model name such as 'resnet-18'
        :param layer: layer as a string for resnet-18 or int for alexnet
        :returns: pytorch model, selected layer
        """
        if model_name == 'resnet-18':
            model = models.resnet18(pretrained = False)
            model.load_state_dict(torch.load(cfg.resnet18_dir))
            if layer == 'default':
                layer = model._modules.get('avgpool')
                self.layer_output_size = 512
            else:
                layer = model._modules.get(layer)

            return model, layer

2.3 方式一加载模型结构的陷阱

需要提前定义模型结构，但是不需要初始化。

或者通过从model文件里引用。from model_save import *

举例，模型保存

# 陷阱
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x
# 定义模型
tudui = Tudui()
torch.save(tudui, "tudui_method1.pth")

模型加载

# 如果不加这段，torch.load会报错
# 陷阱1
# class Tudui(nn.Module):
#     def __init__(self):
#         super(Tudui, self).__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
#
#     def forward(self, x):
#         x = self.conv1(x)
#         return x

model = torch.load('tudui_method1.pth')
print(model)

"""
Tudui(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)
"""

报错信息如下，

    model = torch.load('tudui_method1.pth')
  File "D:\appp\Anaconda3\envs\py37\lib\site-packages\torch\serialization.py", line 387, in load
    return _load(f, map_location, pickle_module, **pickle_load_args)
  File "D:\appp\Anaconda3\envs\py37\lib\site-packages\torch\serialization.py", line 574, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'Tudui' on <module '__main__' from 'C:/Users/。。。/src/model_load.py'>

解决方式，恢复陷阱1的注释 or 加上这段

from model_save import *

注意

如果我不想从torch官方的model进行加载，而是想用我本地已经下载好的model（因为使用官方的model，每次运行调用api都会从网络上下载model，耗时）

# pretrained = False 需要指定False，表示不从网络上下载
model = models.resnet18(pretrained = False)

# 从本地加载resnet18模型，resnet18_dir是模型保存路径
model.load_state_dict(torch.load(resnet18_dir))

sdbhewfoqi

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
1
评论
【Pytorch】网络的保存与读取(6)

一、调用先有的网络模型使用vgg16,https://pytorch.org/vision/stable/models.htmlFalse是初始化的参数，Ture是训练好的参数。如下代码显示了，如何加载pytorch预训练的网络模型，并对网络进行修改。vgg16原始输出结果是1000类，但是CIFAR10只有10类。需要对vgg16进行修改。# -*- coding: utf-8 -*-import torchvision# train_data = torchvis...
复制链接

扫一扫

专栏目录