pytorch并行处理详解

• PyTorch默认使用从0开始的GPU，且默认只使用0号GPU。如果要使用其他编号的GPU或者使用多块GPU，则要设置。

• pytorch并行后，假设batchsize设置为64，表示每张并行使用的GPU都使用batchsize=64来计算（单张卡使用时，使用batchsize=64比较合适时，多张卡并行时，batchsize仍为64比较合适，而不是64*并行卡数）。

• DataParallel 会自动拆分数据，并将作业订单发送到多个GPU上的多个模型。 在每个模型完成它们的工作之后，DataParallel 在将结果返回给你之前收集和合并结果。

有两种方法：

方法一：环境变量 + device + to(device)

第一步，指定备选的GPU

有如下两种方法来指定需要备选的GPU。
这些GPU 将在程序中可见（但不一定会使用，真正投入使用需要通过device()函数来再次选择和指定）

使用环境变量CUDA_VISIBLE_DEVICES的方式，类似于tensorflow指定GPU的方式（http://www.cnblogs.com/darkknightzh/p/6591923.html）。

直接终端中设定：

CUDA_VISIBLE_DEVICES=1 

python代码中设定：

1、使用单卡

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2、使用多块卡的方法。
例如，使用0号和1号GPU’

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'

第二步，创建设备（device）

作用：将备选GPU进一步选择和指定，真正投入使用中。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 注意如果 device = torch.device("cuda")，则环境变量CUDA_VISIBLE_DEVICES中指定的全部GPU都会被拿来使用。
# 也可以通过 "cuda:0" 、"cuda:1"等指定环境变量CUDA_VISIBLE_DEVICES中指定的多块GPU中的某一块。

注意对应关系。例如：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3, 4, 5"  # 将2, 3, 4, 5号GPU作为备选GPU

# 这样cuda:0表示使用 2 号GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

第三步，将data和model放置到device上

# 如果他们已经在目标设备上则不会执行复制操作
input = data.to(device)
model = MyModule(...).to(device)

注意：如果有多个GPU，则model还需要多做一个操作（模型并行化）
第三步骤的多GPU版本如下：

input = data.to(device)
model = MyModule(...)
#模型并行化
if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            model = nn.DataParallel(model)
model = model.to(device)

注意事项

GPU环境变量只在第一次设置起作用，此后再设置不会有任何改变了
运行下面代码，可以观察其作用机制

import os
import torch.backends.cudnn as cudnn

print("The are ", torch.cuda.device_count(), "GPUs in all!")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if not torch.cuda.is_available():
    device = torch.device('cpu')
else:
    cudnn.benchmark = True
    cudnn.enabled = True
    device = torch.device("cuda:1,0")
    print("The environ has", torch.cuda.device_count(), "GPUs!")
print(device)

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")

方法二 函数 set_device + 函数.cuda()

不过官方建议使用CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。

第一步，函数set_device设置device

import torch
gpu_id = [0, 1, 2]
torch.cuda.set_device(gpu_id)

第二部，函数.cuda()使用GPU

data.cuda()
model.cuda()