自用DDP模版
参考:
Pytorch - DDP教程
Pytorch - torchrun 弹性启动
import argparse
import logging
import os
import sys
from importlib import reload
import datasets
import torch
import torch.distributed as dist
import torch.nn as nn
from einops import rearrange
from torch.distributed.elastic.multiprocessing.errors import record
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim.adamw import AdamW
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from torchvision import transforms
class CustomModel(nn.Module):
def __init__(self, dim, heads=4, dim_head=32):
super().__init__()
self.heads = heads
hidden_dim = dim_head * heads
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
def forward(self, x):
b, c, h, w = x.shape
qkv = self.to_qkv(x)
q, k, v = rearrange(
qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
)
k = k.softmax(dim=-1)
context = torch.einsum("bhdn,bhen->bhde", k, v)
out = torch.einsum("bhde,bhdn->bhen", context, q)
out = rearrange(
out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
)
return self.to_out(out)
class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = nn.Linear(10, 10)
self.relu = nn.ReLU()
self.net2 = nn.Linear(10, 5)
def forward(self, x):
return self.net2(self.relu(self.net1(x)))
@record
def main(custom_args):
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
print(f"Start running basic DDP example on rank {rank}.")
# create model and move it to GPU with id rank
gpu = rank % torch.cuda.device_count()
world_size = dist.get_world_size()
torch.cuda.set_device(gpu)
torch.backends.cudnn.benchmark = True
# ...
dataset = Dataset()
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) if (world_size > 1) else None
dataloader = DataLoader(
dataset,
sampler=sampler
)
rank_iter = iter(dataloader)
model=ToyModel().cuda()
ddp_model = DDP(model, device_ids=[gpu])
loss_fn = nn.MSELoss()
optimizer = AdamW(ddp_model.parameters())
optimizer.zero_grad()
# train
# batch = next(rank_iter)
outputs = ddp_model(torch.randn(20, 10))
labels = torch.randn(20, 5).to(gpu)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()
# A step finished
# loss = dist.all_reduce(loss, op=torch.distributed.SUM)
# print(loss / world_size)
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description="Argparser for configuring [code base name to think of] codebase"
)
parser.add_argument("--cfg", type=str, default="config.yaml")
args = parser.parse_args()
main(args)