本記事ではPyTorchを例に、複数GPUを使用した並列学習について説明します。

なぜ複数GPUでの並列学習が必要なのか

その理由は主に2つあります。1つ目はモデルが単一のGPUに収まらない場合、複数のGPU上で完全なモデルを動作させるためです（例えば初期のAlexNetなど）。2つ目は、複数のGPUを並列計算に使用することでトレーニング速度を向上させるためです。

一般的な複数GPUのトレーニング方法

以下の2つのアプローチがあります：

1. モデル分割方式：非常に大きなモデルの場合、メモリ不足を回避するためにネットワークの異なるモジュールをそれぞれ別のGPU上に配置します。
2. データ分割方式：モデル全体を各GPUにコピーし、同時に順伝播と逆伝播を行います。これにより効果的にバッチサイズを拡大できます。

例えば、PyTorch 1.7 + CUDA 10 + Tesla V100環境でResNet34を使用し、バッチサイズ16でSGDアルゴリズムを用いて花のデータセットをトレーニングする場合、1枚のGPUではエポックあたり9秒かかりますが、2枚のGPUでは5.5秒、8枚のGPUでは2秒になります。ここで注意すべき点は、GPU数が増えても通信コストが増加するため、理論的なスケーリング限界があることです。

勾配のデバイス間通信方法

各GPUでステップごとに損失勾配を求め、それらを平均化して同期します。

BN層の同期方法

複数のGPUを使用する際、全デバイスにおける入力データの平均と分散を計算します。これにより、バッチサイズが増加したとみなされ、より良い結果が得られることがあります。

DataParallelとDistributedDataParallelの比較

DataParallelは単一マシン内で動作するシングルプロセス方式ですが、DistributedDataParallelはマルチプロセスに対応しており、単一または複数マシンでの利用が可能です。現在主流なのは後者です。

PyTorchにおけるGPU起動方法

DistributedDataParallelを使用する際の具体的な手順を以下に示します。

1. プロセス環境の初期化

def setup_multiprocessing(args):
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.device_id = int(os.environ['LOCAL_RANK'])
    elif 'SLURM_PROCID' in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.device_id = args.rank % torch.cuda.device_count()
    else:
        print('Distributed mode is not used')
        args.distributed = False
        return

    args.distributed = True
    torch.cuda.set_device(args.device_id)
    dist.init_process_group(backend='nccl', init_method='env://')
    dist.barrier()

2. データセットの準備

データローダーを作成する際に、DistributedSamplerを使用してデータを各GPUに分配します。

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    sampler=train_sampler,
    num_workers=workers,
    pin_memory=True
)

val_loader = torch.utils.data.DataLoader(
    val_dataset,
    batch_size=batch_size,
    sampler=val_sampler,
    num_workers=workers,
    pin_memory=True
)

3. モデルの設定と重みのロード

事前学習済みの重みがあれば、それを読み込みます。

model = custom_model(num_classes=num_classes).to(device)

if os.path.exists(pretrained_path):
    pretrained_dict = torch.load(pretrained_path, map_location=device)
    model.load_state_dict({k: v for k, v in pretrained_dict.items() if model.state_dict()[k].numel() == v.numel()}, strict=False)
else:
    if args.rank == 0:
        torch.save(model.state_dict(), "temp_initial.pth")
    dist.barrier()

    model.load_state_dict(torch.load("temp_initial.pth", map_location=device))

4. 学習ループ

エポックごとにデータサンプラーのエポックを更新し、ランダムなデータシャッフルを確保します。

for epoch in range(total_epochs):
    train_sampler.set_epoch(epoch)
    loss = train_one_epoch(model, optimizer, train_loader, device)
    scheduler.step()

5. 評価フェーズ

評価時には正解数を集計し、それを全てのプロセスで合算します。

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    corrects = torch.zeros(1).to(device)

    for images, labels in loader:
        preds = model(images.to(device))
        preds = torch.argmax(preds, dim=1)
        corrects += torch.eq(preds, labels.to(device)).sum()

    dist.barrier()
    corrects = sum_all(corrects)
    return corrects.item()

def sum_all(tensor):
    dist.all_reduce(tensor)
    return tensor