mnist.patch

--- mnist_main_original.py	2024-08-10 17:30:08.552324326 -0500
+++ pnetcdf_mnist.py	2024-08-11 16:10:31.895471785 -0500
@@ -1,3 +1,8 @@
+#
+# Copyright (C) 2024, Northwestern University and Argonne National Laboratory
+# See COPYRIGHT notice in top-level directory.
+#
+
 import argparse
 import torch
 import torch.nn as nn
@@ -5,7 +10,11 @@
 import torch.optim as optim
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import StepLR
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data.distributed import DistributedSampler

+import comm_file, pnetcdf_io
+from mpi4py import MPI

 class Net(nn.Module):
     def __init__(self):
@@ -42,14 +51,13 @@
         loss = F.nll_loss(output, target)
         loss.backward()
         optimizer.step()
-        if batch_idx % args.log_interval == 0:
+        if rank == 0 and batch_idx % args.log_interval == 0:
             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                 epoch, batch_idx * len(data), len(train_loader.dataset),
                 100. * batch_idx / len(train_loader), loss.item()))
             if args.dry_run:
                 break

-
 def test(model, device, test_loader):
     model.eval()
     test_loss = 0
@@ -62,9 +70,14 @@
             pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
             correct += pred.eq(target.view_as(pred)).sum().item()

+    # aggregate loss among all ranks
+    test_loss = comm.mpi_comm.allreduce(test_loss, op=MPI.SUM)
+    correct = comm.mpi_comm.allreduce(correct, op=MPI.SUM)
+
     test_loss /= len(test_loader.dataset)

-    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+    if rank == 0:
+        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
         test_loss, correct, len(test_loader.dataset),
         100. * correct / len(test_loader.dataset)))

@@ -94,6 +107,8 @@
                         help='how many batches to wait before logging training status')
     parser.add_argument('--save-model', action='store_true', default=False,
                         help='For Saving the current Model')
+    parser.add_argument('--input-file', type=str, required=True,
+                        help='NetCDF file storing train and test samples')
     args = parser.parse_args()
     use_cuda = not args.no_cuda and torch.cuda.is_available()
     use_mps = not args.no_mps and torch.backends.mps.is_available()
@@ -101,18 +116,18 @@
     torch.manual_seed(args.seed)

     if use_cuda:
-        device = torch.device("cuda")
+        torch.cuda.set_device(rank)  # Set the GPU device by rank
+        device = torch.device(f"cuda:{rank}")
     elif use_mps:
         device = torch.device("mps")
     else:
         device = torch.device("cpu")

-    train_kwargs = {'batch_size': args.batch_size}
+    train_kwargs = {'batch_size': args.batch_size//nprocs}
     test_kwargs = {'batch_size': args.test_batch_size}
     if use_cuda:
         cuda_kwargs = {'num_workers': 1,
-                       'pin_memory': True,
-                       'shuffle': True}
+                       'pin_memory': True}
         train_kwargs.update(cuda_kwargs)
         test_kwargs.update(cuda_kwargs)

@@ -120,25 +135,53 @@
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
         ])
-    dataset1 = datasets.MNIST('../data', train=True, download=True,
-                       transform=transform)
-    dataset2 = datasets.MNIST('../data', train=False,
-                       transform=transform)
-    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
-    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+    # Open files storing training and testing samples
+    infile = args.input_file
+    train_file = pnetcdf_io.dataset(infile, 'train_samples', 'train_labels', transform, comm.mpi_comm)
+    test_file = pnetcdf_io.dataset(infile, 'test_samples', 'test_labels', transform, comm.mpi_comm)
+
+    # create distributed samplers
+    train_sampler = DistributedSampler(train_file, num_replicas=nprocs, rank=rank, shuffle=True)
+    test_sampler = DistributedSampler(test_file, num_replicas=nprocs, rank=rank, shuffle=False)
+
+    # add distributed samplers to DataLoaders
+    train_loader = torch.utils.data.DataLoader(train_file, sampler=train_sampler, **train_kwargs)
+    test_loader = torch.utils.data.DataLoader(test_file, sampler=test_sampler, **test_kwargs, drop_last=False)

     model = Net().to(device)
+
+    # use DDP
+    model = DDP(model, device_ids=[device] if use_cuda else None)
+
     optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

     scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
     for epoch in range(1, args.epochs + 1):
+        # train sampler set epoch
+        train_sampler.set_epoch(epoch)
+        test_sampler.set_epoch(epoch)
+
         train(args, model, device, train_loader, optimizer, epoch)
         test(model, device, test_loader)
         scheduler.step()

     if args.save_model:
-        torch.save(model.state_dict(), "mnist_cnn.pt")
+        if rank == 0:
+            torch.save(model.state_dict(), "mnist_cnn.pt")

+    # close files
+    train_file.close()
+    test_file.close()

 if __name__ == '__main__':
+    ## initialize parallel environment
+    comm, device = comm_file.init_parallel()
+
+    rank = comm.get_rank()
+    nprocs = comm.get_size()
+
     main()
+
+    comm.finalize()
+