switch to use PnetCDF-Python to load MNIST data

Author: KWang1998

examples/MNIST/MNIST_codes/main.py

@@ -8,6 +8,40 @@
 import comm_file
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.distributed import ReduceOp, all_reduce
+from pnetcdf import File
+from mpi4py import MPI
+
+class PnetCDFDataset(torch.utils.data.Dataset):
+    def __init__(self, netcdf_file, data_var, label_var, transform=None, comm=MPI.COMM_WORLD):
+        self.netcdf_file = netcdf_file
+        self.data_var = data_var
+        self.label_var = label_var
+        self.transform = transform
+        self.comm = comm
+
+        # Open the NetCDF file
+        self.f = File(self.netcdf_file, mode='r', comm=self.comm)
+        self.f.begin_indep() # To use independent I/O mode
+
+        # Get dimensions of the variables
+        self.data_shape = self.f.variables[self.data_var].shape
+        self.label_shape = self.f.variables[self.label_var].shape
+
+    def __len__(self):
+        return self.data_shape[0]
+
+    def __getitem__(self, idx):
+        # Read the data and label at the given index
+        image = self.f.variables[self.data_var][idx, ...]
+        label = self.f.variables[self.label_var][idx]
+
+        if self.transform:
+            image = self.transform(image)
+
+        return image, label
+
+    def close(self):
+        self.f.close()
 
 class Net(nn.Module):
     def __init__(self):
@@ -118,21 +152,26 @@ def main():
                         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('--netcdf-file', type=str, default="../MNIST_data/mnist_images.nc",
+                        help='netcdf file storing train and test data')
     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()
 
     torch.manual_seed(args.seed)
-
+    
     ## init comm, rank, nprocs
     comm, device = comm_file.init_parallel()
 
     rank = comm.get_rank()
     nprocs = comm.get_size()
+    mpi_comm = MPI.COMM_WORLD
+    mpi_rank = mpi_comm.Get_rank()
+    mpi_size = mpi_comm.Get_size()
 
-    print("nprocs = ", nprocs, " rank = ",rank," device = ", device)
+    print("nprocs = ", nprocs, " rank = ",rank," device = ", device, " mpi_size = ", mpi_size, " mpi_rank = ", mpi_rank)
 
-    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,
@@ -145,11 +184,12 @@ def main():
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
         ])
-    dataset1 = datasets.MNIST('../MNIST_data', train=True, download=True,
-                       transform=transform)
-    dataset2 = datasets.MNIST('../MNIST_data', train=False,
-                       transform=transform)
 
+    # pnetcdf MNIST datasets
+    netcdf_file = args.netcdf_file
+    dataset1 = PnetCDFDataset(netcdf_file, 'train_images', 'train_labels', transform, mpi_comm)
+    dataset2 = PnetCDFDataset(netcdf_file, 'test_images', 'test_labels', transform, mpi_comm)
+ 
     # add train distributed sampler
     train_sampler = torch.utils.data.distributed.DistributedSampler(dataset1, num_replicas=comm.get_size(), rank=comm.get_rank(), shuffle=True)
     test_sampler = torch.utils.data.distributed.DistributedSampler(dataset2, num_replicas=comm.get_size(), rank=comm.get_rank(), shuffle=False)
@@ -175,8 +215,10 @@ def main():
         if rank == 0:
             torch.save(model.state_dict(), "mnist_cnn.pt")
 
+    # close the file
+    dataset1.close()
+    dataset2.close()
     comm.finalize()
 
-
 if __name__ == '__main__':
     main()