[人工智能]网络模型的保存和读取

先新建一个网络模型：

import torchvision

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

pretrained=False意思是该网络模型是初始没有经过训练的。

方式1：

保存数据：

import torch
    
torch.save(vgg16, 'vgg16_method1.pth')

这种方式既保存了网络结构也保存了网络中的参数

加载数据：

import torch

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)

...

正是vgg16的网络结构，通过debug也可以知道二者的网络参数是一样的。

注意：

方式1在load的时候，本.py文件内应该存在该网络，可以通过import导入包含该网络定义的文件，或者可以class定义一个所对应的文件结构。

方式2：

保存数据：

torch.save(vgg16.state_dict(), 'vgg16_method2.pth')

获取vgg16的参数（状态），并转换为字典形式，只保存参数，不保存结构。

加载数据并打印：

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

输出：

OrderedDict([('features.0.weight', tensor([[[[ 4.4961e-02, ?1.9124e-01, ?2.4614e-02],
? ? ? ? ? [-1.7670e-02, -5.0754e-02, -8.5891e-02],
? ? ? ? ? [-1.0293e-01, -5.0360e-02, -7.7070e-02]],

...

输出为字典形式的网络参数，并没有网络结构。若要还原网络，则要创建一个和原模型结构一样的网络，再进行加载：

vgg16 = torchvision.models.vgg16()
vgg16.load_state_dict(torch.load('vgg16_method2.pth'))
print(vgg16)