Merge pull request #6 from Quantum-Accelerators/hanaol/mp-dataloader

Comment about the inefficiency of the dataloading procedure.

Author: hanaol

src/electrai/dataloader/chgcar_read.py

@@ -26,7 +26,6 @@ def __init__(
         train_fraction: fraction of the data used for training (0 to 1).
         '''
         self.data_dir = Path(data_dir)
-        print(self.data_dir)
         self.label_dir = Path(label_dir)
         self.map_dir = Path(map_dir)
         self.rho_type = rho_type

src/electrai/dataloader/dataset.py

@@ -2,6 +2,8 @@
 from torch.utils.data import DataLoader, Dataset
 import numpy as np
 
+# !!! Reading in all the data at once is probably not a good idea, we better just read in the filenames, and have __getitem__ read each index in instead.
+# The could should be updates as such. 
 class RhoData(Dataset):
     def __init__(self, list_data, list_label, list_data_gridsizes, list_label_gridsizes, data_augmentation=True, downsample_data=1, downsample_label=1):
         '''

src/electrai/scripts/train.py

@@ -47,6 +47,8 @@
                     downsample_label=cfg.downsample_label,
                     data_augmentation=False)
 
+print('train_data: ', train_data)
+
 train_loader = DataLoader(train_data, batch_size=int(cfg.nbatch), shuffle=True)
 test_loader = DataLoader(test_data, batch_size=int(cfg.nbatch), shuffle=False)
 
@@ -62,8 +64,6 @@
     normalize=not cfg.normalize_label
 ).to(cfg.device)
 
-print("train chckpt")
-
 optimizer = torch.optim.Adam(model.parameters(), lr=float(cfg.lr), weight_decay=float(cfg.weight_decay))
 
 # Linear + Cosine scheduler
@@ -105,7 +105,13 @@ def loss_fn_sum(output, target):
             return loss
 
     optimizer.zero_grad()
-    for batch, (X, y) in enumerate(dataloader):
+    # print(dataloader)
+    # for batch, (X, y) in enumerate(dataloader):
+    for batch, cont in enumerate(dataloader):
+        print('batch: ', batch)
+        print('cont: ', cont)
+        # print('X: ', X.shape)
+        # print('y: ', y.shape)
         X, y = X.to(cfg.device), y.to(cfg.device)
         pred = model(X)