Pytorch模型的存储与加载

torch.save： 将序列化对象保存到磁盘。此函数使用Python的pickle模块进行序列化。使用此函数可以保存如模型、tensor、字典等各种对象。

torch.load： 使用pickle的unpickling功能将pickle对象文件反序列化到内存。此功能还可以有助于设备加载数据。

torch.nn.Module.load_state_dict： 使用反序列化函数 state_dict 来加载模型的参数字典。

1.state_dict

torch.nn.Module模型的可学习参数（即权重和偏差）包含在模型的参数中，（使用model.parameters()可以进行访问）

print(model.parameters)

state_dict 是Python字典对象，它将每一层映射到其参数张量。注意，只有具有可学习参数的层（如卷积层，线性层等）的模型 才具有 state_dict 这一项。

import torch
import torch.nn as nn
import torch.nn.functional as F
# 定义模型
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class SimpleNet(nn.Module):
    def __init__(self,in_put,hindden,out_put):
        super(SimpleNet,self).__init__()
        self.layers1 = nn.Linear(in_put,hindden)

        self.BatchNormal = nn.BatchNorm1d(hindden)
        self.Relu = nn.ReLU(hindden)  # method one
        self.layers2 = nn.Linear(hindden,hindden)
        self.layers3 = nn.Linear(hindden,out_put)

        self.layers4 = nn.Sequential(
            nn.Linear(in_put,hindden),
            nn.BatchNorm1d(hindden),
            nn.ReLU(True),  #method two
            nn.Linear(hindden,hindden),
            nn.Linear(hindden,out_put)
        )
    def forward(self, x):
        x = self.layers1(x)
        x = self.Relu(x)
        x = self.layers2(x)
        x = self.layers3(x)
        x = self.layers4(x)
        return x
# 初始化模型
model = TheModelClass()

# 初始化优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的状态字典
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 打印优化器的状态字典
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])

Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

2.保存与加载

方式一：保存模型网络结构+参数

torch.save(model,PATH)  #保存

torch.load(PATH)  #加载

方式二：保存网络模型参数

torch.save(model.state_dict(),PATH)

torch.load_state_dict(torch.load(PATH))

load_state_dict()函数只接受字典对象，而不是保存对象的路径。这就意味着在你传给load_state_dict()函数之前，你必须反序列化 你保存的state_dict。例如，你无法通过 model.load_state_dict(PATH)来加载模型。

3.保存和加载 Checkpoint

·保存

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

·加载

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

4.在一个文件中保存多个模型

当保存一个模型由多个 torch.nn.Modules 组成时，例如GAN(对抗生成网络)、sequence-to-sequence (序列到序列模型), 或者是多个模 型融合, 可以采用与保存常规检查点相同的方法。

torch.save({
'modelA_state_dict': modelA.state_dict(),
'modelB_state_dict': modelB.state_dict(),
'optimizerA_state_dict': optimizerA.state_dict(),
'optimizerB_state_dict': optimizerB.state_dict(),
...
}, PATH)

modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)

modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

还有一种方式可以表示

model_encoder = torch(PATH)
model_encoder_dict = model_encoder.state_dict()
====》
model.load_state_dict(model_encoder_dict)

5.使用在不同模型参数下的热启动模式

保存

torch.save(modelA.state_dict(), PATH)

加载

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

无论是从缺少某些键的 state_dict 加载还是从键的数目多于加载模型的 state_dict , 都可以通过在load_state_dict() 函数中将 strict 参数设置为 False 来忽略非匹配键的函数。

6.通过设备保存/加载模型

保存到 CPU，加载到 GPU

# 保存模型
torch.save(model.state_dict(), PATH)

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose
# whatever GPU device number you want
model.to(device)
# 确保在你提供给模型的任何输入张量上调用input = input.to(device)

在CPU上训练好并保存的模型加载到GPU时，将 torch.load() 函数中的 map_location 参数设置为 cuda:device_id 。这会将模型加载到 指定的GPU设备。

接下来，请务必调用 model.to(torch.device(‘cuda’)) 将模型的参数张量转换为 CUDA 张量。最后，确保在所有模型输入上使用 .to(torch.device(‘cuda’)) 函数来为CUDA优化模型。

请注意，调用 my_tensor.to(device) 会在GPU上返回 my_tensor 的新副本。它不会覆盖 my_tensor 。因此， 请手动覆盖张量 my_tensor = my_tensor.to(torch.device(‘cuda’)) 。

数据集类

torch.utils.data.Dataset 是表示数据集的抽象类，因此自定义数据集应继承Dataset并覆盖以下方法 *__len__实现 len(dataset) 返还数据集的尺寸。 *__getitem__用来获取一些索引数据，例如 dataset[i] 中的(i)。

为面部数据集创建一个数据集类。我们将在 ___init___中读取csv的文件内容，在 __getitem__中读取图片。这么做是为了节省内存 空间。只有在需要用到图片的时候才读取它而不是一开始就把图片全部存进内存里。

class FaceLandmarksDataset(Dataset):
	"""面部标记数据集."""
	def __init__(self, csv_file, root_dir, transform=None):
	"""
	csv_file（string）：带注释的csv文件的路径。
	   root_dir（string）：包含所有图像的目录。
	    transform（callable， optional）：一个样本上的可用的可选变换
	"""
		self.landmarks_frame = pd.read_csv(csv_file)
		self.root_dir = root_dir
		self.transform = transform
	def __len__(self):
		return len(self.landmarks_frame)
	def __getitem__(self, idx):
		img_name = os.path.join(self.root_dir,self.landmarks_frame.iloc[idx, 0])
		image = io.imread(img_name)
		landmarks = self.landmarks_frame.iloc[idx, 1:]
		landmarks = np.array([landmarks])
		landmarks = landmarks.astype('float').reshape(-1, 2)
		sample = {'image': image, 'landmarks': landmarks}
		if self.transform:
			sample = self.transform(sample)
		return sample