【Pytorch】pytorch中保存模型的三种方式

1. torch保存模型相关的api

1.1 torch.save()

torch.save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL, _use_new_zipfile_serialization=True)

参考自https://pytorch.org/docs/stable/generated/torch.save.html#torch-save

torch.save()的功能是保存一个序列化的目标到磁盘当中，该函数使用了Python中的pickle库用于序列化，具体参数的解释如下

参数	功能
obj	需要保存的对象
f	指定保存的路径
pickle_module	用于 pickling 元数据和对象的模块
pickle_protocol	指定 pickle protocal 可以覆盖默认参数

常见用法

# dirctly save entiry model
torch.save('model.pth')
# save model'weights only
torch.save(model.state_dict(), 'model_weights.pth')
# save checkpoint
checkpint = {'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': loss,'epoch': epoch
}
torch.save(checkpoint, 'checkpoint_path.pth')

1.2 torch.load()

torch.load(f, map_location=None, pickle_module=pickle, *, weights_only=False, mmap=None, **pickle_load_args)

参考自https://pytorch.org/docs/stable/generated/torch.load.html#torch-load

torch.load()的功能是加载模型，使用python中的unpickle工具来反序列化对象，并且加载到对应的设备上，具体的参数解释如下

参数	功能
f	对象的存放路径
map_location	需要映射到的设备
pickle_module	用于 unpickling 元数据和对象的模块

常见用法

# specify the device to use
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# load entiry model to cuda if available
model = torch.load('whole_model.pth', map_location=device)
# load model's weight to cuda if available
model.load_state_dict(torch.load('model_weights.pth'), map_location=device)
# load checkpoint
checkpoint = torch.load('checkpoint_path.pth', map_location=device)
# checkpoint加载出来就像个字典,预先保存的是否放置了什么内容,加载之后就可以这样来获取
loss = checkpoint['loss']
epoch = chekpoint['epoch']
model.load_state_dict(checkpoint['model_state_dict']
optimizer.load_state_dict(checkpoint['optimizer_state_dict']

1.3 torch.nn.Module.load_state_dict()

torch.nn.Module.load_state_dict(state_dict, strict=True, assign=False)

参考自https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict

torch.nn.Module.load_state_dict()将参数和缓冲区从 state_dict 复制到此模块及其后代中。 如果 strict 为 True，则 state_dict 的键必须与该模块的 state_dict() 函数返回的键完全匹配。具体的参数描述如下

参数	功能
state_dict	保存parameters和persistent buffers的字典
strict	是否强制要求state_dict中的key和model.state_dict返回的key严格一致

1.4 什么是state_dict()

torch.nn.Module.state_dict()

参考自https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.state_dict

其实state_dict可以理解为一种简单的Python Dictionary，其功能是将每层之间的参数进行一一映射并且存储在python的数据类型字典中。因此state_dict可以轻松地进行修改、保存等操作。

除了torch.nn.Module拥有state_dict()方法之外，torch.optim.Optimizer也具有state_dict()方法。如下所示

torch.optim.Optimizer.state_dict()

参考自https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.state_dict.html

1.4. 1 举个例子

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimclass SimpleModel(nn.Module):def __init__(self, input_size, output_size):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(input_size, 100)self.fc2 = nn.Linear(100, output_size)def forward(self, x):x = F.relu(self.fc1(x))return self.fc2(x)if __name__ == "__main__":model = SimpleModel(10, 2)optimizer = optim.Adam(model.parameters(), lr=0.001)print("Check Model's State Dict:")for key, value in model.state_dict().items():print(key, "\t", value.size())print("Check Optimizer's State Dict:")for key, value in optimizer.state_dict().items():print(key, "\t", value)

输出的结果如下

Check Model's State Dict:
fc1.weight       torch.Size([100, 10])
fc1.bias         torch.Size([100])
fc2.weight       torch.Size([2, 100])
fc2.bias         torch.Size([2])
Check Optimizer's State Dict:
state    {}
param_groups     [{'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'differentiable': False, 'fused': None, 'params': [0, 1, 2, 3]}]

---

2. pytorch模型文件后缀

常用的torch模型文件后缀有.pt、.pth，这是最常见的PyTorch模型文件后缀，表示模型的权重、结构和状态字典(state_dict)都被保存在其中。

torch.save(model.state_dict(), 'model_weights.pth')
torch.save(model, 'full_model.pt')

还有检查点后缀如.ckpt、.checkpoint，这些后缀常被用于保存模型的检查点，包括权重和训练状态等。它们也可以表示模型的中间状态，以便在训练期间从中断的地方继续训练。

checkpoint = {'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'epoch': epoch,# 其他信息
}
torch.save(checkpoint, 'model_checkpoint.ckpt')

还有其他跨框架的数据结构例如.h5，PyTorch的模型也可以保存为HDF5文件格式用于跨框架的数据交换，可以使用h5py库来进行读写

import h5pywith h5py.File('model.h5', 'w') as f:# 将模型参数逐一保存到HDF5文件for name, param in model.named_parameters():f.create_dataset(name, data=param.numpy())

3. 存储整个模型

可以直接使用torch.save()和torch.load()来加载和保存整个模型到文件中，这种方式保存了模型的所有权重、架构及其其他相关信息，即使不知道模型的结构也能够直接通过权重文件来加载模型

3.1 直接保存整个模型

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimimport osclass SimpleModel(nn.Module):def __init__(self, input_size, output_size):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(input_size, 256)self.fc2 = nn.Linear(256, 256)self.fc3 = nn.Linear(256, output_size)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))return self.fc3(x)if __name__ == "__main__":model = SimpleModel(10, 2)# specify the save pathurl = os.path.dirname(os.path.realpath(__file__)) + '/models/'# 如果路径不存在则创建if not os.path.exists(url):os.makedirs(url)# specify the model save namemodel_name = 'simple_model.pth'# save the model to filetorch.save(model, url + model_name)

我们直接将模型保存到了当前文件夹下的./models文件夹中，

3.2 直接加载整个模型

由于我们已经保存了模型的所有相关信息，所以我们可以不知道模型的结构也能加载该模型，如下所示

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimimport osclass SimpleModel(nn.Module):def __init__(self, input_size, output_size):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(input_size, 256)self.fc2 = nn.Linear(256, 256)self.fc3 = nn.Linear(256, output_size)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))return self.fc3(x)if __name__ == "__main__":device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# model = SimpleModel(10, 2)# specify the save pathurl = os.path.dirname(os.path.realpath(__file__)) + '/models/'# 如果路径不存在则创建if not os.path.exists(url):os.makedirs(url)# specify the model save namemodel_name = 'simple_model.pth'# load the modelif os.path.exists(url + model_name):model = torch.load(url + model_name, map_location=device)print("Success Load Model From:\n\t%s"%(url+model_name))

成功加载了模型

---

4. 只保存模型的权重

4.1 保存模型权重

利用前面提到的state_dict()方法来完成这一操作

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimimport osclass SimpleModel(nn.Module):def __init__(self, input_size, output_size):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(input_size, 256)self.fc2 = nn.Linear(256, 256)self.fc3 = nn.Linear(256, output_size)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))return self.fc3(x)if __name__ == "__main__":# specify devicedevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')model = SimpleModel(10, 2)# specify the save pathurl = os.path.dirname(os.path.realpath(__file__)) + '/models/'# 如果路径不存在则创建if not os.path.exists(url):os.makedirs(url)# specify the model save namemodel_name = 'simple_model_weights.pth'torch.save(model.state_dict(), url + model_name)

我们直接将模型权重保存到了当前文件夹下的./models文件夹中，

4.2 读取模型权重

由于我们只保存了模型的权重信息，不知道模型的结构，所以必须要先实例化模型才行。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimimport osclass SimpleModel(nn.Module):def __init__(self, input_size, output_size):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(input_size, 256)self.fc2 = nn.Linear(256, 256)self.fc3 = nn.Linear(256, output_size)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))return self.fc3(x)if __name__ == "__main__":# specify devicedevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# get modelmodel = SimpleModel(10, 2)# specify the save pathurl = os.path.dirname(os.path.realpath(__file__)) + '/models/'# 如果路径不存在则创建if not os.path.exists(url):os.makedirs(url)# specify the model save namemodel_name = 'simple_model_weights.pth'if os.path.exists(url + model_name):model.load_state_dict(torch.load(url + model_name, map_location=device))print("Success Load Model'weights From:\n\t%s"%(url+model_name))

5. 使用Checkpoint保存中间结果

5.1 保存Checkpoint

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import os# 数据准备
x = torch.tensor(np.random.rand(100, 1), dtype=torch.float32)
y = 3 * x + 2 + 0.1 * torch.randn(100, 1)# 定义模型
class SimpleLinearModel(nn.Module):def __init__(self):super(SimpleLinearModel, self).__init__()self.linear = nn.Linear(1, 1)def forward(self, x):return self.linear(x)if __name__=="__main__":# specify devicedevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# 实例化模型model = SimpleLinearModel()# 定义损失函数和优化器criterion = nn.MSELoss()optimizer = optim.SGD(model.parameters(), lr=0.01)# 训练循环num_epochs = 1000checkpoint_interval = 100  # 保存检查点的间隔url = os.path.dirname(os.path.realpath(__file__))+'/models/'if not os.path.exists(url):os.makedirs(url)checkpoint_file = 'checkpoint.pth'  # 检查点文件路径for epoch in range(num_epochs):# 前向传播outputs = model(x)loss = criterion(outputs, y)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()# 打印训练信息if (epoch + 1) % checkpoint_interval == 0:print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')# 保存检查点checkpoint = {'epoch': epoch + 1,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'loss': loss.item(),}torch.save(checkpoint, url+checkpoint_file)

5.2 加载Checkpoint

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import os# 数据准备
x = torch.tensor(np.random.rand(100, 1), dtype=torch.float32)
y = 3 * x + 2 + 0.1 * torch.randn(100, 1)# 定义模型
class SimpleLinearModel(nn.Module):def __init__(self):super(SimpleLinearModel, self).__init__()self.linear = nn.Linear(1, 1)def forward(self, x):return self.linear(x)if __name__=="__main__":# specify devicedevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')# 实例化模型model = SimpleLinearModel()# 定义损失函数和优化器criterion = nn.MSELoss()optimizer = optim.SGD(model.parameters(), lr=0.01)# 训练循环num_epochs = 1000checkpoint_interval = 100  # 保存检查点的间隔url = os.path.dirname(os.path.realpath(__file__))+'/models/'if not os.path.exists(url):os.makedirs(url)checkpoint_file = 'checkpoint.pth'  # 检查点文件路径# load from checkpointcheckpoint = torch.load(url+checkpoint_file)for key, value in checkpoint.items():print(key, '-->', value)model.load_state_dict(checkpoint['model_state_dict'])optimizer.load_state_dict(checkpoint['optimizer_state_dict'])epoch = checkpoint['epoch']loss = checkpoint['loss']print('Loaded checkpoint from epoch %d. Loss %f' % (epoch, loss))

输出如下

loss --> 0.01629752665758133
(test_ros_python) sjh@sjhR9000X:~/Documents/python_draft$  cd /home/sjh/Documents/python_draft ; /usr/bin/env /home/sjh/anaconda3/envs/metaRL/bin/python /home/sjh/.vscode/extensions/ms-python.python-2023.18.0/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher 40897 -- /home/sjh/Documents/python_draft/check_checkpoint.py 
epoch --> 1000
model_state_dict --> OrderedDict([('linear.weight', tensor([[2.6938]])), ('linear.bias', tensor([2.1635]))])
optimizer_state_dict --> {'state': {0: {'momentum_buffer': None}, 1: {'momentum_buffer': None}}, 'param_groups': [{'lr': 0.01, 'momentum': 0, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False, 'params': [0, 1]}]}
loss --> 0.01629752665758133
Loaded checkpoint from epoch 1000. Loss 0.016298

我们成功从断点处加载checkpoint, 可以再从这个断点处继续训练

Reference

参考一