Update mnist_main.py patch

Pytorch github repo has updated mnist_main.py on Apr 30, 2025, which
uses torch.accelerator APIs. See
pytorch/examples#1334

Author: wkliao

examples/MNIST/mnist.patch

@@ -1,8 +1,8 @@
---- mnist_main_original.py	2024-08-10 17:30:08.552324326 -0500
-+++ pnetcdf_mnist.py	2024-08-11 16:10:31.895471785 -0500
+--- mnist_main.py.orig	2025-05-09 10:51:06.814200110 -0500
++++ mnist_main.py	2025-05-09 11:15:17.198167820 -0500
 @@ -1,3 +1,8 @@
 +#
-+# Copyright (C) 2024, Northwestern University and Argonne National Laboratory
++# Copyright (C) 2025, Northwestern University and Argonne National Laboratory
 +# See COPYRIGHT notice in top-level directory.
 +#
 +
@@ -15,13 +15,13 @@
  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 @@
+@@ -42,7 +51,7 @@
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
@@ -30,62 +30,47 @@
              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 @@
+@@ -62,9 +71,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 @@
+ 
+@@ -92,6 +106,8 @@
                          help='how many batches to wait before logging training status')
-     parser.add_argument('--save-model', action='store_true', default=False,
+     parser.add_argument('--save-model', action='store_true', 
                          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")
+ 
+     use_accel = not args.no_accel and torch.accelerator.is_available()
+@@ -103,12 +119,11 @@
      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,
+     if use_accel:
+         accel_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 @@
+         train_kwargs.update(accel_kwargs)
+         test_kwargs.update(accel_kwargs)
+ 
+@@ -116,25 +131,53 @@
          transforms.ToTensor(),
          transforms.Normalize((0.1307,), (0.3081,))
          ])
@@ -108,14 +93,14 @@
 +    # 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)
++    model = DDP(model, device_ids=[device] if use_accel 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
@@ -125,16 +110,16 @@
          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()