pytorch保存及加载模型的两种方式

import os.path
from typing import Iterator

import numpy as np
import torch
import cv2
from PIL import Image
from torch.utils.data import Dataset,DataLoader
import re
from functools import reduce
from torch.utils.tensorboard import SummaryWriter as Writer
from torchvision import transforms
import torchvision as tv
from torch import nn
import torch.nn.functional as F
#
t1=transforms.ToTensor()
#中心裁剪
#传入tensor，输出tensor；传入PIL，输出PIL
t3=transforms.Resize((500,500))
t=transforms.Compose([t1,t3])
#If train=True, creates dataset from training set, otherwise creates from test set.
#数据集下载后为PIL格式：
myDataSet=tv.datasets.CIFAR10(root="./myDataSet",train=False,download=True,transform=t)
dataloader=DataLoader(myDataSet,shuffle=True,batch_size=32,num_workers=0,drop_last=False)
#查看命令：tensorboard --logdir=./myBorderText
#可用pycharm中code中的generater功能实现：
class myCustomerNetWork(nn.Module):
    def __init__(self):
        super().__init__()
        #输入3通道输出6通道：
        self.conv1=nn.Conv2d(3,6,(10,10))
    def forward(self,x):
        x=nn.ReLU()(self.conv1(x))
        x=nn.MaxPool2d(5)(x)
        x=nn.Conv2d(6,3,(5,5))(x)
        x = nn.ReLU()(x)
        x=nn.MaxPool2d(5)(x)
        x = nn.ReLU()(self.conv1(x))
        x=nn.Flatten()(x)
        x=nn.Linear(x.shape[1],10)(x)
        return x

myNet=myCustomerNetWork()
#模型保存，同时保存模型和参数：
torch.save(myNet,"./saveText.pth")
#模型加载：加载自定义模型时必须在该py文件中存在自定义类
restoredNet=torch.load("./saveText.pth")
print(restoredNet)
#或者保存模型参数，而不保存模型,参数以python字典的形式输出：
torch.save(myNet.state_dict(),"./saveTextOnlyParams.pth")
#模型加载相关参数：
myNet.load_state_dict(torch.load("./saveTextOnlyParams.pth"))