Pytorch：模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load |
您所在的位置：网站首页 › archive和save › Pytorch：模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load |
|
Pytorch 保存和加载模型后缀：.pt 和.pth 1 torch.save() [source]保存一个序列化（serialized）的目标到磁盘。函数使用了Python的pickle程序用于序列化。模型（models），张量（tensors）和文件夹（dictionaries）都是可以用这个函数保存的目标类型。 torch.save(obj, f, pickle_module=, pickle_protocol=2) 参数描述obj保存对象f类文件对象 (必须实现写和刷新)或一个保存文件名的字符串pickle_modul用于 pickling 元数据和对象的模块 pickle_protocol指定 pickle protocal 可以覆盖默认参数 栗子保存整个模型： torch.save(model,'save.pt')只保存训练好的权重： torch.save(model.state_dict(), 'save.pt') 2 torch.load() [source]用来加载模型。torch.load() 使用 Python 的 解压工具（unpickling）来反序列化 pickled object 到对应存储设备上。首先在 CPU 上对压缩对象进行反序列化并且移动到它们保存的存储设备上，如果失败了（如：由于系统中没有相应的存储设备），就会抛出一个异常。用户可以通过 register_package 进行扩展，使用自己定义的标记和反序列化方法。 torch.load(f, map_location=None, pickle_module=) 参数描述f类文件对象 (返回文件描述符)或一个保存文件名的字符串map_location一个函数或字典规定如何映射存储设备pickle_module用于 unpickling 元数据和对象的模块 (必须匹配序列化文件时的 pickle_module ) 栗子 torch.load('tensors.pt') # Load all tensors onto the CPU torch.load('tensors.pt', map_location=torch.device('cpu')) # Load all tensors onto the CPU, using a function torch.load('tensors.pt', map_location=lambda storage, loc: storage) # Load all tensors onto GPU 1 torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1)) # Map tensors from GPU 1 to GPU 0 torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'}) # Load tensor from io.BytesIO object with open('tensor.pt') as f: buffer = io.BytesIO(f.read()) torch.load(buffer) 3 torch.nn.Module.load_state_dict(state_dict) [source]使用 state_dict 反序列化模型参数字典。用来加载模型参数。将 state_dict 中的 parameters 和 buffers 复制到此 module 及其子节点中。 torch.nn.Module.load_state_dict(state_dict, strict=True) 参数描述state_dict保存 parameters 和 persistent buffers 的字典strict可选，bool型。state_dict 中的 key 是否和 model.state_dict() 返回的 key 一致。 栗子 torch.save(model,'save.pt') model.load_state_dict(torch.load("save.pt")) #model.load_state_dict()函数把加载的权重复制到模型的权重中去 3.1 什么是state_dict？在PyTorch中，一个torch.nn.Module模型中的可学习参数（比如weights和biases），模型的参数通过model.parameters()获取。而state_dict就是一个简单的Python dictionary，其功能是将每层与层的参数张量之间一一映射。注意，只有包含了可学习参数（卷积层、线性层等）的层和已注册的命令（registered buffers，比如batchnorm的running_mean）才有模型的state_dict入口。优化方法目标（torch.optim）也有state_dict，其中包含的是关于优化器状态的信息和使用到的超参数。 因为state_dict目标是Python dictionaries，所以它们可以很轻松地实现保存、更新、变化和再存储，从而给PyTorch模型和优化器增加了大量的模块化（modularity）。 torch.nn.Module.state_dict torch.nn.Module.state_dict(destination=None, prefix='', keep_vars=False)返回一个包含模型状态信息的字典。包含参数（weighs and biases）和持续的缓冲值（如：观测值的平均值）。只有具有可更新参数的层才会被保存在模型的 state_dict 数据结构中。 栗子： module.state_dict().keys() # ['bias', 'weight']torch.optim.Optimizer.state_dict torch.optim.Optimizer.state_dict()返回一个包含优化器状态信息的字典。包含两个 key： state：字典，保存当前优化器的状态信息。不同优化器内容不同。param_groups：字典，包含所有参数组（eg：超参数）。 栗子： from __future__ import print_function, division import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import numpy as np import torchvision from torchvision import datasets, models, transforms import matplotlib.pyplot as plt import time import os import copy # 定义模型 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 # 初始化模型 model = TheModelClass() # 初始化优化器 optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 打印模型的 state_dict print("Model's state_dict:") for param_tensor in model.state_dict(): print(param_tensor, "\t", model.state_dict()[param_tensor].size()) # 打印优化器的 state_dict 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]}] 4 保存/加载 4.1 state_dict（推荐） 保存： torch.save(model.state_dict(), PATH) 加载： model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.eval()保存模型的推理过程的时候，只需要保存模型训练好的参数，使用torch.save()保存state_dict，能够方便模型的加载。因此推荐使用这种方式进行模型保存。 记住一定要使用model.eval()来固定dropout和归一化层，否则每次推理会生成不同的结果。 注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict() 4.2 整个模型 保存 torch.save(model, PATH) 加载 # 模型类必须在别的地方定义 model = torch.load(PATH) model.eval()这种保存/加载模型的过程使用了最直观的语法，所用代码量少。这使用Python的pickle保存所有模块。这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。 一般来说，PyTorch的模型以.pt或者.pth文件格式保存。 一定要记住在评估模式的时候调用model.eval()来固定dropout和批次归一化。否则会产生不一致的推理结果。 4.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() # - 或者 - model.train()在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。 要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。 加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。 同样，评估模型的时候一定不要忘了调用model.eval()。 4.4 保存多个模型到一个文件 保存 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']) modelA.eval() modelB.eval() # - 或者 - modelA.train() modelB.train()保存的模型包含多个torch.nn.Modules时，比如GAN，一个序列-序列模型，或者组合模型，使用与保存常规检查点的方式来保存模型。也就是说，保存每个模型的state_dict和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。 4.5 使用其他模型来预热当前模型 保存 torch.save(modelA.state_dict(), PATH) 加载 modelB = TheModelBClass(*args, **kwargs) modelB.load_state_dict(torch.load(PATH), strict=False)在迁移学习或者训练新的复杂模型时，加载部分模型是很常见的。利用经过训练的参数，即使只有少数参数可用，也将有助于预热训练过程，并且使模型更快收敛。 在加载部分模型参数进行预训练的时候，很可能会碰到键不匹配的情况（模型权重都是按键值对的形式保存并加载回来的）。因此，无论是缺少键还是多出键的情况，都可以通过在load_state_dict()函数中设定strict参数为False来忽略不匹配的键。 如果想将某一层的参数加载到其他层，但是有些键不匹配，那么修改state_dict中参数的key可以解决这个问题。 4.6 跨设备保存与加载模型 GPU上保存，CPU上加载 保存 torch.save(model.state_dict(), PATH) 加载 device = torch.device('cpu') model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location=device))当在CPU上加载一个GPU上训练的模型时，在torch.load()中指定map_location=torch.device('cpu')，此时，map_location动态地将tensors的底层存储重新映射到CPU设备上。 上述代码只有在模型是在一块GPU上训练时才有效，如果模型在多个GPU上训练，那么在CPU上加载时，会得到类似如下错误： KeyError: ‘unexpected key “module.conv1.weight” in state_dict’原因是在使用多GPU训练并保存模型时，模型的参数名都带上了module前缀，因此可以在加载模型时，把key中的这个前缀去掉： # 原始通过DataParallel保存的文件 state_dict = torch.load('myfile.pth.tar') # 创建一个不包含`module.`的新OrderedDict from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # 去掉 `module.` new_state_dict[name] = v # 加载参数 model.load_state_dict(new_state_dict)GPU上保存，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")) # 选择希望使用的GPU model.to(device)在把GPU上训练的模型加载到GPU上时，只需要使用model.to(torch.devie('cuda'))将初始化的模型转换为CUDA优化模型。同时确保在模型所有的输入上使用.to(torch.device('cuda'))。注意，调用my_tensor.to(device)会返回一份在GPU上的my_tensor的拷贝。不会覆盖原本的my_tensor，因此要记得手动将tensor重写：my_tensor = my_tensor.to(torch.device('cuda'))。 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) # Make sure to call input = input.to(device) on any input tensors that you feed to the model在 GPU 上加载 CPU 训练保存的模型时，将 torch.load() 函数的 map_location 参数 设置为 cuda:device_id。这种方式将模型加载到指定设备。下一步，确保调用 model.to(torch.device(‘cuda’)) 将模型参数 tensor 转换为 cuda tensor。最后，确保模型输入使用 .to(torch.device(‘cuda’)) 为 cuda 优化模型准备数据。 注意：调用 my_tensor.to(device) 会在 GPU 上返回 my_tensor 的新副本，不会覆盖 my_tensor。因此，使用 my_tensor = my_tensor.to(torch.device(‘cuda’)) 手动覆盖。 4.7 保存torch.nn.DataParallel模型 保存 torch.save(model.state_dict(), PATH)加载 # Load to whatever device you wanttorch.nn.DataParallel 是支持模型使用 GPU 并行的封装器。要保存一个一般的 DataParallel 模型， 请保存 model.module.state_dict()。这种方式，可以灵活地以任何方式加载模型到任何设备上。 5 参考 https://blog.csdn.net/LXYTSOS/article/details/90639524https://blog.csdn.net/cpongo3/article/details/93624288 |
| CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3 |