PyTorch深度学习（10）不同网络模型及网络模型保存、加载

一、不同的网络模型

CIFAR10、VGG16、ImageNet等

导入库

pip install scipy    安装scipy

pip uninstall scipy  卸载scipy

torchvision.datasets.ImageNet('./dataset', split='train', download=True, transform=torchvision.transforms.ToTensor())

参数：存放位置、 split——训练集、是否下载、transform转换

VGG16

150000张图片，1000中类别，神经网络经过30步

outfeatures = 1000

具体代码

import torchvision
from torch import nn

# 报错，The dataset is no longer publicly accessible 不能公开下载，只能下载后替换跟目录文件
# 此文件train 训练集图片大小为147.9GB
# torchvision.datasets.ImageNet('./dataset', split='train', download=True,
#                               transform=torchvision.transforms.ToTensor())

# 数据量为528M
# 网络模型参数是未经过训练的,初始化的参数
vgg16_false = torchvision.models.vgg16(pretrained=False)
# 网络模型参数是经过训练的  
vgg16_true = torchvision.models.vgg16(pretrained=True)

print(vgg16_true)

train_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=True,
                             transform=torchvision.transforms.ToTensor())

# 向vgg16中添加模块 使得由1000类输出为10类
vgg16_true.classifier.add_module('add_linear', nn.Linear(1000, 10))
print(vgg16_true)

print(vgg16_false)
# 指定步骤由4096类输出10类
vgg16_false.classifier[6] = nn.Linear(4096, 10)
print(vgg16_false)

二、网络模型的保存

torch.save(数据集，文件名称)    不仅保存了网络模型的结构，还保存了网络模型的参数

torch.save(vgg16.state_dict()，文件名称)  ——官方推荐

保存的是vgg16网络参数保存为字典形式 

具体代码

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)

# 保存方式1，模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")

# 保存方式2，模型参数（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

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

    def forward(self, input):
        output = self.conv1(input)
        return output


t1 = Test1()
torch.save(t1, "test_method1")

三、网络模型的读取

torch.load(文件名)

具体代码

import torch
import torchvision
from torch import nn

# 或者解决方法  将model_save中的代码引入
from model_save import *

# 方式一保存后的模型加载
model = torch.load("vgg16_method1.pth")
# print(model)

# 方式二加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
model = vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# model = torch.load("vgg16_method2.pth")
# 输出为字典
# print(model)

# 陷阱1，
# 解决
# class Test1(nn.Module):
#     def __init__(self):
#         super(Test1, self).__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
#
#     def forward(self, input):
#         output = self.conv1(input)
#         return output


model = torch.load("test_method1")  # 报错AttributeError: Can't get attribute 'Test1' on <module '__main__'
print(model)