PyTorch中如何实现GPU与GPU、GPU与CPU间的数据加载转换操作？

本文共计454个文字，预计阅读时间需要2分钟。

问题描述：有时在加载已训练好的模型时，会出现out of memory的错误提示，尽管进行了详细的内存检查，使用的GPU卡并没有再使用，内存也未超出限制。

经查，发现此问题原因为训练模型时使用的GPU卡和加载模型时使用的GPU卡不一致。

问题描述

有时在加载已训练好的模型时，会出现 out of memory 的错误提示，但仔细检测使用的GPU卡并没有再用且内存也没有超出。

经查阅发现原来是训练模型时使用的GPU卡和加载时使用的GPU卡不一样导致的。个人感觉，因为pytorch的模型中是会记录有GPU信息的，所以有时使用不同的GPU加载时会报错。

解决方法

gpu之间的相互转换。即，将训练时的gpu卡转换为加载时的gpu卡。

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'}) # gpu 1 --> gpu 0

当然，你也可以直接将加载模型时使用的gpu卡改为和训练时的一样。但在多人使用一个服务器时，你想用的gpu卡已被使用，就需按上面方法转换gpu。