pytorch在gpu训练加载模型,训练与测试在不同机器上遇到的问题总结

今天遇到了在测试的时候遇到了一件有趣的事情：
模型是我在自己电脑上训练出来的，cuda:0，测试的时候为了加快程序运行的速度弄到了实验室的服务器上去跑，服务器一共两个GPU，0和1，0被另一个同学占用了，所以我通过device = torch.device(“cuda:1” if torch.cuda.is_available() else “cpu”)去指定指用GPU1，(我怀疑我没有指定上，因为我俩的程序都跑了一会之后自动停了，error是Process finished with exit code 137 (interrupted by signal 9: SIGKILL)，显然是内存不足了。) 后来我又通过torch.cuda.set_device(1)指定，通过torch.cuda.current_device() 返回的设备索引这次应该是指定上了。but，，，问题来了，在checkpoint=torch.load(“model.pth”)，预加载模型时竟然老是给我报错：CUDA:out of memory！ 但是通过nvidia-smi查看GPU的使用率时显示1号GPU的使用率并不高啊，我就纳闷了，然后让同学把0号GPU的程序停了，我用0号GPU跑试了试，加载模型完全没有问题！！！
找原因，突然想到模型在训练的时候是在我的电脑上GPU为0， 当我在服务器上指定用1号GPU上加载模型时就总是出错，原因可能是模型的参数是cuda:0, 而我在服务器的1号GPU上加载就不行？？？于是鼓捣了一下：checkpoint=torch.load(‘model.pth’, map_location={‘cuda:0’:‘cuda:1’})把张量从GPU 0 移动到GPU 1，程序顺利运行。。。。。
看来pytorch对设备把关挺严。

再加上今天学习到的一个新的小知识：在测试时，也可以不把model放到gpu上，即使加载的模型是在gpu上训练的，我们只load模型的参数，在output=model(input),时，input也不放到GPU上就不会报错，但是，如果前面有model.to(device),后面input没有放到GPU上，就会报错，即为model与input的device要统一。
当然还有多GPU训练，单GPU测试的，这种的遇到的时候会再补充学习的。

仅此Mark