【Pytorch1.0.0系列】RuntimeError: arguments are located on different GPUs at ...

报错信息

RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorIndex.cu:519

Pytorch1.0.0当中的多GPU使用

# 绝对误差损失函数测试
import os
import re
import torch
import torch.nn as nn
#实现例如： python filename.py --gpu_id=0,1,2,3,4,5,6,7 的参数传递
parser = argparse.ArgumentParser(description='manual')
parser.add_argument('--gpu_id', help="GPU_ID", type=str, default = "0,1,2,3,4")
USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda:"+re.split(r",",args.gpu_id)[0] if USE_CUDA else "cpu")
gpu_id = list(map(int, re.split(r",",args.gpu_id)))
# 需要强调的是，所有.cuda()均需要用.to(device)来替代，否则在运行程序的时候就会报错：
# RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorIndex.cu:519
class yourmodel(nn.Module):
    def __init__(self,
                 args1 = default1,
                 argsn = defaultn):
                 ...
    ...
    def train()
    	...
    def predict()
    	...

    ...
if __name__ == '__main__':
    model = yourmodel()
    
    model = torch.nn.DataParallel(model, device_ids = gpu_id)
    model.to(device)
    if args.run_type == 'train':
        #seq.train()  # 单GPU
        model.module.train()   # 加上.module
    elif args.run_type == 'predict':
        #seq.predict()  # 单GPU
        model.module.predict()   # 加上.module

如果出现报错RuntimeError: arguments are located on different GPUs at

可能的原因1：

在网络外要用到网络中的子模块，却没有加上.module，例如：

if __name__ == '__main__':
    model = yourmodel()
    
    model = torch.nn.DataParallel(model, device_ids = gpu_id)
    model.to(device)
    if args.run_type == 'train':
        seq.train()  # 单GPU
    elif args.run_type == 'predict':
        seq.predict()  # 单GPU

正确的做法：

if __name__ == '__main__':
    model = yourmodel()
    
    model = torch.nn.DataParallel(model, device_ids = gpu_id)
    model.to(device)
    if args.run_type == 'train':
        seq.module.train()  # 多GPU
    elif args.run_type == 'predict':
        seq.module.predict()  # 多GPU

可能的原因2：

某一处没有使用".to(device)"，例如：

        if self.use_cuda:
            decoder_input = decoder_input.cuda()
            decoder_context = decoder_context.cuda()

正确的做法：

        if self.use_cuda:
            decoder_input = decoder_input.to(device)
            decoder_context = decoder_context.to(device)