pytorch单机多卡训练--完整示例代码

指定可用GPU

在多人共用多卡的情况下，并不是任何时刻每张卡都是空闲的，因此，在模型训练和推理时，需要指定可用的gpu。pytorch中推荐的方法为，使用 os.environ["CUDA_VISIBLE_DEVICES"] 设置可用的gpu。

import torch
# import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

print(torch.cuda.device_count())
print(torch.cuda.current_device())

我的服务器有4张卡，运行上述代码后，运行结果为 4，0。当前设备的默认索引为0；

取消上述注释代码，重新运行，运行结果为2，0。因为 os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" 指定了可用的gpu索引，因此可用gpu的数量变为2。

若一个服务器上有多个GPU设备，可以使用多GPU设备进行训练，充分利用多GPU计算的性能，缩短训练时长。

在具有多GPU的服务器上，通常可以有两种提高训练效率的方式：

• 数据并行（DataParallel）
• 分布式数据并行（DistributedDataParallel）

数据并行（DataParallel）

在训练时我们通常会把训练样本划分为 mini-batches，一次进行一个 batch size 的计算；多GPU上的数据并行是指将一个 batch size 的计算平均地分配到多个GPU上进行并行计算，从而提高计算效率。

例如，一个 batch_size = 32，可用于计算的 GPU 数量为4，一个 batch_size 的计算将分配到4个 GPU 上并行计算，一个 GPU 一次计算8个样本的数据。

在 pytorch 中使用 [torch.nn.DataParallel](https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html?highlight=torch%20nn%20dataparallel#torch.nn.DataParallel) 实现多GPU的并行计算，一个简单的可运行案例如下，

参考链接：

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 32
data_size = 128

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)

class Model(nn.Module):
    # 自己定义的模型

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())

        return output

model = Model(input_size, output_size)

# GPU数量大于1才能实现多GPU并行计算
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [32, xxx] -> [8, ...], [8, ...], [8, ...] on 4 GPUs
  # 使用 DataParallel 把模型包装起来
  model = nn.DataParallel(model)

# 将模型送到多个GPU上
model.to(device)

for data in rand_loader:
    # 将数据送到多个GPU上
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

我实验的服务器上有4块卡，运行结果如下：

DataParallel 会自动拆分输入的数据并将工作指令发送到多个 GPU 上的多个模型。在每个模型完成其工作后，DataParallel 会收集并合并结果，然后再将其返回。

pytorch将数据并行计算的方法包装的非常方便，只需在原来代码的基础上加上一两行代码就可以实现，但它通常不能提供最佳性能，因为它在每个前向传递中复制模型，并且其并行模式为单进程多线程，因此自然会受到 GIL 争用的影响。

分布式数据并行（DistributedDataParallel，DDP）

DDP使用多进程并行，并且模型在DDP构建时就进行复制。

以单机多GPU为例。

1. 构建并初始化进程组

DDP使用多进程进行并行计算，在模型、数据等配置准备前，需要先构建并初始化DDP的运行环境。

[torch.distributed.init_process_group()](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) 函数提供了初始化DDP运行环境的功能。在 pytorch 的官网中，介绍了三种初始化方法，这里只介绍使用环境变量进行初始化的方式。

我们主要关心以下四个环境变量：

• MASTER_ADDR - 节点为0的ip地址，因为这里是单机，因此设置为 localhost；
• MASTER_PORT - ****标识为0的进程上的可用端口；
• WORLD_SIZE - 进程总数，可以通过环境变量设置，也可以在 init_process_group 函数中设置；
• RANK - 进程标识，用于设置当前进程标识，相当于设置主进程的标识；可以通过环境变量设置，也可以在 init_process_group 函数中设置；

示例如下：

def ddp_setup(rank, world_size):
    """
    Args:
        rank: 进程的唯一标识，在 init_process_group 中用于指定当前进程标识
        world_size: 进程总数
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

2. 分发数据

与使用单机单卡不同，单机多卡需要将数据分发到多个gpu上，我们只需要在 [torch.utils.data.DataLoader](https://pytorch.org/docs/stable/data.html?highlight=torch+utils+data+dataloader#torch.utils.data.DataLoader) 中指定 sampler 即可实现，代码示例如下：

dataloader = torch.utils.data.DataLoader(
    dataset=train_dataset,
    batch_size=32,
    shuffle=False,
    sampler=torch.utils.data.distributed.DistributedSampler(train_dataset),
)

如上所示，每个gpu的一个 batch size 为32，如果gpu的可用数量为4，那么有效的 batch size 为 32 x 4；

需要注意的是，在单机单卡训练中，我们一般会将 shuffle 设置为 True，以使得每个 epoch 训练中模型看见数据的顺序都不一样。但在多卡训练中，因为指定了 sampler，我们将 shuffle 设置为 False。DistributedSampler 实现了多卡训练中数据的 shuffle 功能，此外，为了保证每个epoch训练中，数据shuffle成功，需要在每个 epoch 的 dataloader 前调用 set_epoch() 方法，示例如下：

for epoch in range(max_epochs):
		dataloader.sampler.set_epoch(epoch)
		for features, labels in dataloader:
				...

3. 构建DDP模型

初始化DDP的运行环境后就可以构建DDP模型了，示例如下：

from torch.nn.parallel import DistributedDataParallel as DDP

model = DDP(model, device_ids=[gpu_id])      # gpu_id 为我们传递的 rank 参数

4. 保存模型

因为每个进程保存的模型都是相同的，因此我们只需要保存一个进程中的模型权重。示例如下：

checkpoint = self.model.module.state_dict()
# 保存进程标识为0下的模型权重
if self.gpu_id == 0 and epoch % self.save_every == 0:
		save_checkpoint(epoch, checkpoint)

5. 开始运行多卡训练

与单机单卡训练不同，使用 torch.multiprocessing.spawn 启动多卡训练的主程序，spawn 会自动分配 rank 参数。

对于表示进程数的 world_size 参数，应该和当前设备中的可用GPU数量保持一致，即 world_size = torch.cuda.device_count()

待补充：在linux下，python中使用多进程，好像进程使用 spawn 方式创建进程，那 spawn 到底是个啥？

6. 一个单机多卡训练的完整运行示例

import torch
from torch.utils.data import Dataset
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
import os

class MyTrainDataset(Dataset):
    def __init__(self, size):
        self.size = size
        self.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)]

    def __len__(self):
        return self.size
    
    def __getitem__(self, index):
        return self.data[index]

def ddp_setup(rank, world_size):
    """
    Args:
        rank: 进程的唯一标识，在 init_process_group 中用于指定当前进程标识
        world_size: 进程总数
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    init_process_group(backend="nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        gpu_id: int,
        save_every: int,
    ) -> None:
        self.gpu_id = gpu_id
        self.model = model.to(gpu_id)
        self.train_data = train_data
        self.optimizer = optimizer
        self.save_every = save_every
        self.model = DDP(model, device_ids=[gpu_id])

    def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.model(source)
        loss = F.cross_entropy(output, targets)
        loss.backward()
        self.optimizer.step()

    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")
        self.train_data.sampler.set_epoch(epoch)
        for source, targets in self.train_data:
            source = source.to(self.gpu_id)
            targets = targets.to(self.gpu_id)
            self._run_batch(source, targets)

    def _save_checkpoint(self, epoch):
        ckp = self.model.module.state_dict()
        PATH = "checkpoint.pt"
        torch.save(ckp, PATH)
        print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")

    def train(self, max_epochs: int):
        for epoch in range(max_epochs):
            self._run_epoch(epoch)
            if self.gpu_id == 0 and epoch % self.save_every == 0:
                self._save_checkpoint(epoch)

def load_train_objs():
    train_set = MyTrainDataset(2048)  # load your dataset
    model = torch.nn.Linear(20, 1)  # load your model
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
    return train_set, model, optimizer

def prepare_dataloader(dataset: Dataset, batch_size: int):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,
        sampler=DistributedSampler(dataset)
    )

def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):
    ddp_setup(rank, world_size)
    dataset, model, optimizer = load_train_objs()
    train_data = prepare_dataloader(dataset, batch_size)
    trainer = Trainer(model, train_data, optimizer, rank, save_every)
    trainer.train(total_epochs)
    destroy_process_group()

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='simple distributed training job')
    parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')
    parser.add_argument('save_every', type=int, help='How often to save a snapshot')
    parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)')
    args = parser.parse_args()
    
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size)

待补充运行截图…

使用torchrun进行多卡训练（待更新…）

参考资料

https://pytorch.org/tutorials/beginner/ddp_series_intro.html

https://github.com/pytorch/examples/tree/main/distributed/ddp-tutorial-series

https://www.youtube.com/watch?v=-LAtx9Q6DA8