How to perform GPU-to-GPU and GPU-to-CPU data conversion when loading in PyTorch.
发布时间
阅读量:
阅读量
pytorch中 如何将gpu与gpu、gpu与cpu 在load时相互转化载入
在实际应用过程中,有时需要将原本在CPU上训练的模型迁移至GPU进行处理,或因特定条件限制,将原本在GPU上训练的模型转移回CPU运行。当在GPU环境下完成模型训练后,如何高效地将模型权重迁移到CPU设备中?本文将对此问题展开探讨:
将模型部署至目标设备
根据官方技术文档提供的信息,我们可以参考以下实现方式
torch.load('tensors.pt')
# 把所有的张量加载到CPU中
torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# 把所有的张量加载到GPU 1中
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# 把张量从GPU 1 移动到 GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
将预先训练完成的GPU模型部署至CPU环境时,存在一种强制性处理方式,即要求所有原本
全部评论 (0)
还没有任何评论哟~
