Advertisement

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)

还没有任何评论哟~