00081-PyTorch 数据并行

PyTorch
PyTorch
发布日期:   2023-07-22
更新日期:   2024-12-28

前言

PyTorch 的官网地址为：https://pytorch.org/ 。

PyTorch Tutorials 的地址为：https://pytorch.org/tutorials/ 。

操作系统：Windows 10 专业版

参考文档

  1. MULTI GPU TRAINING WITH DDP

安装 PyTorch

源教程地址: https://pytorch.org/tutorials/beginner/ddp_series_multigpu.html .

官方代码仓库: https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu.py .

  1. 新建 multigpu.py 文件:
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

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


def ddp_setup(rank, world_size):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    init_process_group(backend="gloo", 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)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# Download training data from open datasets.
train_set = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

def load_train_objs():
    model = NeuralNetwork()
    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()
    print(f"total_epochs: {args.total_epochs}, save_every: {args.save_every}, batch_size: {args.batch_size}")
    
    world_size = torch.cuda.device_count()
    mp.spawn(main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size)
  1. 运行代码:
$ python multigpu.py --help
usage: multigpu.py [-h] [--batch_size BATCH_SIZE] total_epochs save_every

simple distributed training job

positional arguments:
  total_epochs          Total epochs to train the model
  save_every            How often to save a snapshot

options:
  -h, --help            show this help message and exit
  --batch_size BATCH_SIZE
                        Input batch size on each device (default: 32)
$ python multigpu.py 50 10
total_epochs: 50, save_every: 10, batch_size: 32
[GPU0] Epoch 0 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 0 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 1 | Batchsize: 32 | Steps: 938
Epoch 0 | Training checkpoint saved at checkpoint.pt
[GPU0] Epoch 1 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 2 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 2 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 3 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 3 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 4 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 4 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 5 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 5 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 6 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 6 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 7 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 7 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 8 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 8 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 9 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 9 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 10 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 10 | Batchsize: 32 | Steps: 938
Epoch 10 | Training checkpoint saved at checkpoint.pt
[GPU1] Epoch 11 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 11 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 12 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 12 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 13 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 13 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 14 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 14 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 15 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 15 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 16 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 16 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 17 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 17 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 18 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 18 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 19 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 19 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 20 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 20 | Batchsize: 32 | Steps: 938
Epoch 20 | Training checkpoint saved at checkpoint.pt
[GPU1] Epoch 21 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 21 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 22 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 22 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 23 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 23 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 24 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 24 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 25 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 25 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 26 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 26 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 27 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 27 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 28 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 28 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 29 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 29 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 30 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 30 | Batchsize: 32 | Steps: 938
Epoch 30 | Training checkpoint saved at checkpoint.pt
[GPU1] Epoch 31 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 31 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 32 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 32 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 33 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 33 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 34 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 34 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 35 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 35 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 36 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 36 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 37 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 37 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 38 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 38 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 39 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 39 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 40 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 40 | Batchsize: 32 | Steps: 938
Epoch 40 | Training checkpoint saved at checkpoint.pt
[GPU1] Epoch 41 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 41 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 42 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 42 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 43 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 43 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 44 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 44 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 45 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 45 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 46 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 46 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 47 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 47 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 48 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 48 | Batchsize: 32 | Steps: 938
[GPU1] Epoch 49 | Batchsize: 32 | Steps: 938
[GPU0] Epoch 49 | Batchsize: 32 | Steps: 938
$

其他官方教程

  1. https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
  2. https://pytorch.org/docs/stable/notes/ddp.html

结语

第八十一篇博文写完，开心！！！！

今天，也是充满希望的一天。

文章作者: LuYF-Lemon-love
文章链接: https://luyf-lemon-love.space/2023/07/22/00081-pytorch-shu-ju-bing-xing/
版权声明: 本博客所有文章除特別声明外，均采用 CC BY 4.0 许可协议。转载请注明来源 LuYF-Lemon-love !
PyTorch
赏
  目录