torch.nn.DataParallel使用细节

import torch
from torch import nn
#使用机器有两个GPU

a=torch.rand([5,4])
model=nn.Sequential(
      nn.Linear(4,2)
)

model=nn.DataParallel(model)
print(model.device_ids)
b=model(a)
print(b)

#报错，RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu
#解决方案：model=nn.DataParallel(model).to(torch.device('cuda'))
#或者：
#model.to(torch.device('cuda'))
#model=nn.DataParallel(model)

#The parallelized module must have its parameters and buffers on device_ids[0] before running this DataParallel module.
#在运行此DataParallel模块之前，并行化模块的参数和缓冲区必须位于device_ids[0]上。
#也就是说model=nn.DataParallel(model).to(torch.device('cuda:1'))
#会报错RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1

"""
解决报错后运行结果：
[0, 1]
tensor([[-0.7412, -0.2679],
        [-0.4051, -0.2946],
        [-0.5498, -0.2944],
        [-0.3615, -0.2638],
        [-0.2613, -0.2737]], device='cuda:0', grad_fn=<GatherBackward>)
"""

#注意：
#>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
#>>> output = net(input_var)  # input_var can be on any device, including CPU

#torch.device('cuda')默认为torch.device('cuda:0')

当cuda：0显存不够时，可以使用例如：os.environ["CUDA_VISIBLE_DEVICES"] = "3,2,0,1"

这样的命令来把cuda：3修改为cuda：0。