PyTorch(15)---模型保存和加载

模型保存和加载

模型保存

import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear


#模型：vgg16
vgg16=torchvision.models.vgg16(pretrained=False)
'''第一种:模型参数都保存'''
torch.save(vgg16,"vgg16_method1.pth")
'''第二种:只保存参数，保存成字典（官方推荐）'''
torch.save(vgg16.state_dict(),"vgg16_method2.pth")


'''陷阱1:使用自己搭建的网络(第一种)'''
class Demo(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model1=Sequential(
            Conv2d(in_channels=3, out_channels=32, kernel_size=5, stride=1, padding=2, dilation=1, ),
            MaxPool2d(kernel_size=2, ),
            Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2, ),
            MaxPool2d(kernel_size=2),
            Conv2d(32, 64, 5, 1, 2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10),
        )
    def forward(self,x):
        x=self.model1(x)
        return x
#模型2
demo=Demo()
torch.save(demo,"demo_method1.pth")

模型加载

import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from model_save import *

'''方式1:(使用保存方式1)，加载模型'''
model=torch.load("vgg16_method1.pth")
print(model)
print("------------------------------------------------------------")
'''方式2'''
#先把结构搞出来
vgg16=torchvision.models.vgg16(pretrained=False)
#再把字典形式的参数放进去(自己训练的参数！！！)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)
print("------------------------------------------------------------")


'''陷阱1:要把自己的网络模型放过来，但是不用实例化了 或者将模型保存的文件import过来 '''
medel1=torch.load("demo_method1.pth")
print(medel1)