Improving performance of data loading and collecting

dlrm_data_pytorch.py

@@ -220,9 +220,9 @@ def __init__(
                     indices = np.random.permutation(indices)
                     print("Randomized indices...")
 
-                X_int[indices] = X_int
-                X_cat[indices] = X_cat
-                y[indices] = y
+                self.X_int = X_int[indices]
+                self.X_cat = X_cat[indices]
+                self.y = y[indices]
 
             else:
                 indices = np.array_split(indices, self.offset_per_file[1:-1])
@@ -246,17 +246,17 @@ def __init__(
 
                 # create training, validation, and test sets
                 if split == 'train':
-                    self.X_int = [X_int[i] for i in train_indices]
-                    self.X_cat = [X_cat[i] for i in train_indices]
-                    self.y = [y[i] for i in train_indices]
+                    self.X_int = X_int[train_indices]
+                    self.X_cat = X_cat[train_indices]
+                    self.y = y[train_indices]
                 elif split == 'val':
-                    self.X_int = [X_int[i] for i in val_indices]
-                    self.X_cat = [X_cat[i] for i in val_indices]
-                    self.y = [y[i] for i in val_indices]
+                    self.X_int = X_int[val_indices]
+                    self.X_cat = X_cat[val_indices]
+                    self.y = y[val_indices]
                 elif split == 'test':
-                    self.X_int = [X_int[i] for i in test_indices]
-                    self.X_cat = [X_cat[i] for i in test_indices]
-                    self.y = [y[i] for i in test_indices]
+                    self.X_int = X_int[test_indices]
+                    self.X_cat = X_cat[test_indices]
+                    self.y = y[test_indices]
 
             print("Split data according to indices...")
 
@@ -328,9 +328,9 @@ def __len__(self):
 def collate_wrapper_criteo_offset(list_of_tuples):
     # where each tuple is (X_int, X_cat, y)
     transposed_data = list(zip(*list_of_tuples))
-    X_int = torch.log(torch.tensor(transposed_data[0], dtype=torch.float) + 1)
-    X_cat = torch.tensor(transposed_data[1], dtype=torch.long)
-    T = torch.tensor(transposed_data[2], dtype=torch.float32).view(-1, 1)
+    X_int = torch.log(torch.tensor(np.array(transposed_data[0]), dtype=torch.float) + 1)
+    X_cat = torch.tensor(np.array(transposed_data[1]), dtype=torch.long)
+    T = torch.tensor(np.array(transposed_data[2]), dtype=torch.float32).view(-1, 1)
 
     batchSize = X_cat.shape[0]
     featureCnt = X_cat.shape[1]
@@ -399,9 +399,9 @@ def diff(tensor):
 def collate_wrapper_criteo_length(list_of_tuples):
     # where each tuple is (X_int, X_cat, y)
     transposed_data = list(zip(*list_of_tuples))
-    X_int = torch.log(torch.tensor(transposed_data[0], dtype=torch.float) + 1)
-    X_cat = torch.tensor(transposed_data[1], dtype=torch.long)
-    T = torch.tensor(transposed_data[2], dtype=torch.float32).view(-1, 1)
+    X_int = torch.log(torch.tensor(np.array(transposed_data[0]), dtype=torch.float) + 1)
+    X_cat = torch.tensor(np.array(transposed_data[1]), dtype=torch.long)
+    T = torch.tensor(np.array(transposed_data[2]), dtype=torch.float32).view(-1, 1)
 
     batchSize = X_cat.shape[0]
     featureCnt = X_cat.shape[1]

dlrm_s_pytorch.py

@@ -887,7 +887,7 @@ def inference(
             ),
             flush=True,
         )
-    return model_metrics_dict, is_best
+    return model_metrics_dict, is_best, best_acc_test
 
 
 def run():
@@ -1658,7 +1658,7 @@ def run():
                         print(
                             "Testing at - {}/{} of epoch {},".format(j + 1, nbatches, k)
                         )
-                        model_metrics_dict, is_best = inference(
+                        model_metrics_dict, is_best, best_acc_test = inference(
                             args,
                             dlrm,
                             best_acc_test,