01.16 pytorch学习（网络模型的保存与读取）

总所周知我们写的网络可以写在一个python文件中，通过包导入。但也可以手动保存

torch提供了save的方法可以保存网络
两种保存方式
1.存取网络结构和参数
2.存取网络参数

保存方式

import torch
import torchvision
import torch.nn as nn

#True为加载已经训练的参数
vgg16 = torchvision.models.vgg16(pretrained=True)

#已有网络
#保存方式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,32,5,padding=2)

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

tudui = Tudui()
torch.save(tudui,"tudui_method1.pth")

读取方式：

import torch
import torchvision
import torch.nn as nn

model1 = torch.load("vgg16_method1.pth")
print(model1)#输出的网络结构

model2 = torch.load("vgg16_method2.pth")
print(model2)#输出的网络参数
#怎么加载参数到网络结构中呢
#False为不加载已经训练的参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)#输出的网络结构

#tudui = torch.load("tudui_method1.pth")
#print(tudui)
#AttributeError: Can't get attribute 'Tudui' on <module '__main__' from 'C:/Users/2477491565/Desktop/learn_pytorch/model_load.py'>
#需要将这个模块添加到此文件中
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui,self).__init__()
        self.conv1 = nn.Conv2d(3,32,5,padding=2)

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

tudui = torch.load("tudui_method1.pth")
print(tudui)