Add distributed training examples of PyTorch

examples/Distributed-example/Lite-imagenet-singal-mul-DDP-nccl-gloo.py

@@ -0,0 +1,117 @@
+import os
+from datetime import datetime
+import argparse
+import torch.multiprocessing as mp
+import torch.backends.cudnn as cudnn
+import torchvision
+import torchvision.transforms as transforms
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from apex.parallel import DistributedDataParallel as DDP
+from apex import amp
+
+import torchvision.datasets as datasets
+import torchvision.models as models
+model_names = sorted(name for name in models.__dict__
+    if name.islower() and not name.startswith("__")
+    and callable(models.__dict__[name]))
+def main():
+    print('run main')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('data', metavar='DIR',
+                        help='path to dataset')
+    parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
+                        choices=model_names,
+                        help='model architecture: ' +
+                             ' | '.join(model_names) +
+                             ' (default: resnet18)')
+    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N',
+                        help='number of data loading workers (default: 4)')
+    parser.add_argument('-g', '--gpus', default=1, type=int,
+                        help='number of gpus per node')
+    parser.add_argument('-nr', '--nr', default=0, type=int,
+                        help='ranking within the nodes')
+    parser.add_argument('-b', '--batch-size', default=256, type=int,
+                        metavar='N',
+                        help='mini-batch size (default: 256), this is the total '
+                             'batch size of all GPUs on the current node when '
+                             'using Data Parallel or Distributed Data Parallel')
+    parser.add_argument('--epochs', default=2, type=int, metavar='N',
+                        help='number of total epochs to run')
+    parser.add_argument('--dist-backend',  default='nccl', type=str,
+                        help='distributed backend')
+    args = parser.parse_args()
+    args.world_size = args.gpus * args.nodes
+    print('world_size:',args.world_size)
+    os.environ['MASTER_ADDR'] = os.environ['PAI_HOST_IP_worker_0']
+    os.environ['MASTER_PORT'] = os.environ['PAI_worker_0_SynPort_PORT']
+    print('master:', os.environ['MASTER_ADDR'], 'port:', os.environ['MASTER_PORT'])
+    mp.spawn(train, nprocs=args.gpus, args=(args,))
+
+def train(gpu, args):
+    print("start train")
+    rank = int(os.environ['PAI_TASK_INDEX']) * args.gpus + gpu
+    dist.init_process_group(backend=args.dist_backend, init_method='env://', world_size=args.world_size, rank=rank)
+    torch.manual_seed(0)
+    model=model = models.__dict__[args.arch]()
+    torch.cuda.set_device(gpu)
+    model.cuda(gpu)
+    batch_size = args.batch_size
+    # define loss function (criterion) and optimizer
+    criterion = nn.CrossEntropyLoss().cuda(gpu)
+    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
+    # Wrap the model
+    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
+    # Data loading code
+    traindir = os.path.join(args.data, 'train')
+    valdir = os.path.join(args.data, 'val')
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+    train_dataset = datasets.ImageFolder(
+        traindir,
+        transforms.Compose([
+            transforms.RandomResizedCrop(224),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.nodes, pin_memory=True, sampler=train_sampler)
+
+    val_loader = torch.utils.data.DataLoader(
+        datasets.ImageFolder(valdir, transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            normalize,
+        ])),
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.nodes, pin_memory=True)
+    start = datetime.now()
+    total_step = len(train_loader)
+    for epoch in range(args.epochs):
+        for i, (images, labels) in enumerate(train_loader):
+            images = images.cuda(non_blocking=True)
+            labels = labels.cuda(non_blocking=True)
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            # Backward and optimize
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            #if (i + 1) % 100 == 0 and gpu == 0:
+            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
+                                                                         loss.item()))
+    if gpu == 0:
+        print("Training complete in: " + str(datetime.now() - start))
+
+
+if __name__ == '__main__':
+    main()

examples/Distributed-example/LiteApex-imagenet-singal-mul-DDP-nccl-gloo.py

@@ -0,0 +1,122 @@
+import os
+from datetime import datetime
+import argparse
+import torch.multiprocessing as mp
+import torch.backends.cudnn as cudnn
+import torchvision
+import torchvision.transforms as transforms
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from apex.parallel import DistributedDataParallel as DDP
+from apex import amp
+
+import torchvision.datasets as datasets
+import torchvision.models as models
+model_names = sorted(name for name in models.__dict__
+    if name.islower() and not name.startswith("__")
+    and callable(models.__dict__[name]))
+def main():
+    print('run main')
+    parser = argparse.ArgumentParser()
+    parser.add_argument('data', metavar='DIR',
+                        help='path to dataset')
+    parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
+                        choices=model_names,
+                        help='model architecture: ' +
+                             ' | '.join(model_names) +
+                             ' (default: resnet18)')
+    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N',
+                        help='number of data loading workers (default: 4)')
+    parser.add_argument('-g', '--gpus', default=1, type=int,
+                        help='number of gpus per node')
+    parser.add_argument('-nr', '--nr', default=0, type=int,
+                        help='ranking within the nodes')
+    parser.add_argument('-b', '--batch-size', default=256, type=int,
+                        metavar='N',
+                        help='mini-batch size (default: 256), this is the total '
+                             'batch size of all GPUs on the current node when '
+                             'using Data Parallel or Distributed Data Parallel')
+    parser.add_argument('--epochs', default=2, type=int, metavar='N',
+                        help='number of total epochs to run')
+    parser.add_argument('--dist-backend',  default='nccl', type=str,
+                        help='distributed backend')
+    args = parser.parse_args()
+    args.world_size = args.gpus * args.nodes
+    print('world_size:',args.world_size)
+    os.environ['MASTER_ADDR'] = os.environ['PAI_HOST_IP_worker_0']
+    os.environ['MASTER_PORT'] = os.environ['PAI_worker_0_SynPort_PORT']
+    print('master:', os.environ['MASTER_ADDR'], 'port:', os.environ['MASTER_PORT'])
+    mp.spawn(train, nprocs=args.gpus, args=(args,))
+
+def train(gpu, args):
+    print("start train")
+    rank = int(os.environ['PAI_TASK_INDEX']) * args.gpus + gpu
+    dist.init_process_group(backend=args.dist_backend, init_method='env://', world_size=args.world_size, rank=rank)
+    torch.manual_seed(0)
+    model=model = models.__dict__[args.arch]()
+    torch.cuda.set_device(gpu)
+    model.cuda(gpu)
+    batch_size = args.batch_size
+    # define loss function (criterion) and optimizer
+    criterion = nn.CrossEntropyLoss().cuda(gpu)
+    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
+    # Wrap the model
+    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
+    # Wrap the model
+    model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
+    model = DDP(model)
+    # Data loading code
+    traindir = os.path.join(args.data, 'train')
+    valdir = os.path.join(args.data, 'val')
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+
+    train_dataset = datasets.ImageFolder(
+        traindir,
+        transforms.Compose([
+            transforms.RandomResizedCrop(224),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.nodes, pin_memory=True, sampler=train_sampler)
+
+    val_loader = torch.utils.data.DataLoader(
+        datasets.ImageFolder(valdir, transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            normalize,
+        ])),
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.nodes, pin_memory=True)
+    start = datetime.now()
+    total_step = len(train_loader)
+    for epoch in range(args.epochs):
+        for i, (images, labels) in enumerate(train_loader):
+            images = images.cuda(non_blocking=True)
+            labels = labels.cuda(non_blocking=True)
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            # Backward and optimize
+            optimizer.zero_grad()
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            loss.backward()
+            optimizer.step()
+            #if (i + 1) % 100 == 0 and gpu == 0:
+            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
+                                                                         loss.item()))
+    if gpu == 0:
+        print("Training complete in: " + str(datetime.now() - start))
+
+if __name__ == '__main__':
+    main()