PyTorch模型保存与加载的方法

这篇文章主要介绍了PyTorch模型保存与加载的方法的相关知识，内容详细易懂，操作简单快捷，具有一定借鉴价值，相信大家阅读完这篇PyTorch模型保存与加载的方法文章都会有所收获，下面我们一起来看看吧。输出：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]}]保存：加载：要注意这个细节，如果使用nn.DataParallel在一台电脑上使用了多个GPU，那么加载模型的时候也必须先进行nn.DataParallel。保存模型的推理过程的时候，只需要保存模型训练好的参数，使用torch.save()保存state_dict，能够方便模型的加载。因此推荐使用这种方式进行模型保存。记住一定要使用model.eval()来固定dropout和归一化层，否则每次推理会生成不同的结果。注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict()保存：加载：这种保存/加载模型的过程使用了最直观的语法，所用代码量少。这使用Python的pickle保存所有模块。这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。一般来说，PyTorch的模型以.pt或者.pth文件格式保存。一定要记住在评估模式的时候调用model.eval()来固定dropout和批次归一化。否则会产生不一致的推理结果。保存：加载：在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。同样，评估模型的时候一定不要忘了调用model.eval()。保存：加载：保存的模型包含多个torch.nn.Modules时，比如GAN，一个序列-序列模型，或者组合模型，使用与保存常规检查点的方式来保存模型。也就是说，保存每个模型的state_dict和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。保存：加载：在迁移学习或者训练新的复杂模型时，加载部分模型是很常见的。利用经过训练的参数，即使只有少数参数可用，也将有助于预热训练过程，并且使模型更快收敛。在加载部分模型参数进行预训练的时候，很可能会碰到键不匹配的情况（模型权重都是按键值对的形式保存并加载回来的）。因此，无论是缺少键还是多出键的情况，都可以通过在load_state_dict()函数中设定strict参数为False来忽略不匹配的键。如果想将某一层的参数加载到其他层，但是有些键不匹配，那么修改state_dict中参数的key可以解决这个问题。GPU上保存，CPU上加载保存：加载：当在CPU上加载一个GPU上训练的模型时，在torch.load()中指定map_location=torch.device(‘cpu’)，此时，map_location动态地将tensors的底层存储重新映射到CPU设备上。上述代码只有在模型是在一块GPU上训练时才有效，如果模型在多个GPU上训练，那么在CPU上加载时，会得到类似如下错误：原因是在使用多GPU训练并保存模型时，模型的参数名都带上了module前缀，因此可以在加载模型时，把key中的这个前缀去掉：GPU上保存，GPU上加载保存：加载：在把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.nn.DataParallel模型保存：关于“PyTorch模型保存与加载的方法”这篇文章的内容就介绍到这里，感谢各位的阅读！相信大家对“PyTorch模型保存与加载的方法”知识都有一定的了解，大家如果还想学习更多知识，欢迎关注百云主机行业资讯频道。

相关推荐: go-zero单体服务使用泛型简化注册Handler路由的问题怎么解决

这篇文章主要讲解了“go-zero单体服务使用泛型简化注册Handler路由的问题怎么解决”，文中的讲解内容简单清晰，易于学习与理解，下面请大家跟着小编的思路慢慢深入，一起来研究和学习“go-zero单体服务使用泛型简化注册Handler路由的问题怎么解决”吧…