Pytorch 读取参数错误 RuntimeError: cuda runtime error (10) : invalid device ordinal

很久没发博客了，但是今天在Pytorch的参数读取过程中遇到了一个比较罕见的bug。

RuntimeError: cuda runtime error (10) : invalid device ordinal at torch/csrc/cuda/Module.cpp:87

找了很久没有解决方案，StackOverflow和pytorch的issues上也没有人遇到相似的问题。最后只能亲亲自上阵，逐步调试的时候从底层看到了问题，于是在这里记录一下希望能帮到他人。

1. 问题场景
我在服务器上训练了一个模型，把它保存的参数拿到本地机上load进行结果分析的时候，发生了这个错误。在Load参数的时候，提示我cuda找不到了。详细点的报错信息如下：

THCudaCheck FAIL file=torch/csrc/cuda/Module.cpp line=87 error=10 : invalid device ordinal
Traceback (most recent call last):
  File "/home/sw/Shin/Codes/DL4SS_Keras/Torch_multi/main_run_multi_selfSS_subeval.py", line 557, in <module>
    main()
  File "/home/sw/Shin/Codes/DL4SS_Keras/Torch_multi/main_run_multi_selfSS_subeval.py", line 482, in main
    mix_hidden_layer_3d.load_state_dict(torch.load('params/param_mix101_dbag1nosum_WSJ0_hidden3d_190'))
  File "/usr/local/lib/python2.7/dist-packages/torch/serialization.py", line 231, in load
    return _load(f, map_location, pickle_module)
  File "/usr/local/lib/python2.7/dist-packages/torch/serialization.py", line 379, in _load
    result = unpickler.load()
  File "/usr/local/lib/python2.7/dist-packages/torch/serialization.py", line 350, in persistent_load
    data_type(size), location)
  File "/usr/local/lib/python2.7/dist-packages/torch/serialization.py", line 85, in default_restore_location
    result = fn(storage, location)
  File "/usr/local/lib/python2.7/dist-packages/torch/serialization.py", line 67, in _cuda_deserialize
    return obj.cuda(device_id)
  File "/usr/local/lib/python2.7/dist-packages/torch/_utils.py", line 58, in _cuda
    with torch.cuda.device(device):
  File "/usr/local/lib/python2.7/dist-packages/torch/cuda/__init__.py", line 128, in __enter__
    torch._C._cuda_setDevice(self.idx)
RuntimeError: cuda runtime error (10) : invalid device ordinal at torch/csrc/cuda/Module.cpp:87

我本来以为是我的本地机CUDA出了问题，但是如果换一个之前保存的参数，竟然就没有这个问题，很正常的可以work。分析下来，应该就是这个要读取的参数本身有些问题。

底层查看之后，发现了问题。原来是Pytorch在参数保存的时候，会注册一个跟原来参数位置有关的location。比如原来你在服务器上的GPU1训练，这个location很可能就是GPU1了。而如果你台式机上只有一个GPU，也就是GPU0的时候，那么这个参数带进来的Location信息于你的台式机不兼容，就会发生找不到cuda device的问题了。

2.解决方案
在Load参数的时候，根据你现在机器上的GPU状态进行map。例如原来的GPU1可以转换成GPU0，这个在torch.load函数里有个参数可以选。

load(f, map_location=None, pickle_module=pickle)

在我的实验中，最终就是这样一句：

att_speech_layer.load_state_dict(torch.load('params—xxxxx',map_location={'cuda:1':'cuda:0'}))

就解决了，当然得根据你的电脑调整一下这个map。

Pytorch 读取参数错误 RuntimeError: cuda runtime error (10) : invalid device ordinal

猜你喜欢