本文解决了一种常见的PyTorch多GPU模型预测加载错误问题，该问题出现在从多GPU训练的模型切换到单GPU预测时。通过调整代码中CUDA_VISIBLE_DEVICES环境变量和DataParallel的device_ids参数，成功解决了CUDA error invalid device ordinal的运行时错误。
跑了两天的模型，拿出来做预测的时候加载出了问题，报错如下：
| RuntimeError: CUDA error (10): invalid device ordinal (check_status at /opt/conda/conda-bld/pytorch_1532579805626/work/aten/src/ATen/cuda/detail/CUDAHooks.cpp:36) |
最后发现是因为自己训练的时候用了多卡：
就是训练代码里的这两行代码惹的祸
|
os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1, 2, 3" net_model = nn.DataParallel(net_model, device_ids=[0, 1, 2, 3]) |
针对这种情况，解决方法如下：
将预测代码里的这两行改为如下：
|
os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1, 2, 3" net_model = nn.DataParallel(net_model, device_ids=[0]) |
扫一扫
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
