pytorch-多GPU训练（单机多卡、多机多卡）

pytorch 单机多卡训练

首先是数据集的分布处理

需要用到的包：
torch.utils.data.distributed.DistributedSampler
torch.utils.data.DataLoader
torch.utils.data.Dataset

DistributedSampler这个包我们用来确保dataloader只会load到整个数据集的一个特定子集，为每一个子进程划分出一部分数据集，以避免不同进程之间数据重复。

而我们构造的Sampler是在构造DataLoader时的一个参数

先构建一个自己的torch.utils.data.Dataset类的dataset类:

import torch.utils.data as Data
class MyDataSet(Data.Dataset):def __init__(self, x_train, x_label, y_train, y_label): super(MyDataSet, self).__init__()self.x_train = x_trainself.x_label = x_labelself.y_train = y_trainself.y_label = y_labeldef __len__(self):return self.x_train.shape[0]def __getitem__(self, idx):return self.x_train[idx], self.x_label[idx], self.y_train[idx], self.y_label[idx]

初始化DataSet:
这里的x_train, x_label, y_train, y_label是数据集加载进来的变量，根据自己的数据集情况设置变量和数量

train_data = MyDataSet(x_train, x_label, y_train, y_label)

然后根据构造的Dataset对象构造一个DistributedSampler：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_data, shuffle=True)

然后使用torch.utils.data.DataLoader构造自己的DataLoader

MyDataLoader = Data.DataLoader(MyDataSet(x_train, x_label, \y_train, y_label), batch_size=2, sampler = train_sampler)

然后是模型的分布处理

需要用到的包：
torch.nn.parallel.DistributedDataParallel
在初始化model之后，进行一步：

 model = torch.nn.parallel.DistributedDataParallel(model)# device_ids will include all GPU devices by default

pytorch 多机多卡训练

多机多卡需要配置多机间通讯，进行多进程的初始化：
需要用到的包：
torch.distributed
在初始化model之后，进行一步：

torch.distributed.init_process_group(init_method='tcp://{}:{}'.format(args.ip, args.port),backend="gloo",

参数解析：

backend: 多个机器之间交换数据的协议
init_method: 机器之间交换数据, 需要指定一个主节点，这代表主节点
world_size: 进程个数,实际上就是机器的个数, 比如四台台机器一起训练的话, world_size就为4
rank: 区分主节点和从节点的, 主节点为0, 剩余的为了1-(N-1), N为要使用的机器的数量, 也就是world_size