网络模型的保存和读取

1. 网络保存 

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')

#保存自定义的模型
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3,64,5)

    def forward(self,x):
        x = self.conv1(x)
        return x

tudui = Tudui()

torch.save(tudui,'tudui_method.pt')

2. 网络加载

import torch

# 加载方式1 --保存方式1
model1 = torch.load('vgg16_method1.pth')
print(model1)
#
# 加载方式2 --保存方式2
model2 = torch.load('vgg16_method2.pth')
print(model2) #打印出的是参数

#如何获取模型结构？
import torchvision

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

#加载自定义保存的模型
from model_save import tudui
tudui = torch.load('tudui_method.pt')
print(tudui)  #失败的话需要引入包