Merge pull request #4 from danissomo/main

RAM consumption fix and the quality of life changes

Author: Anton Andreychuk

.gitignore

@@ -127,4 +127,8 @@ venv.bak/
 dmypy.json
 
 # Pyre type checker
-.pyre/
+.pyre/
+
+
+TransPath_data
+wandb

data/hmaps.py

@@ -108,26 +108,27 @@ class GridData(Dataset):
         cf - correction factor values
     """
     def __init__(self, path, mode='f', clip_value=0.95):
-        maps = np.load(os.path.join(path, 'maps.npy')).astype('float32')
-        self.maps = torch.tensor(maps)
-        goals = np.load(os.path.join(path, 'goals.npy')).astype('float32')
-        self.goals = torch.tensor(goals)
-        starts = np.load(os.path.join(path, 'starts.npy')).astype('float32')
-        self.starts = torch.tensor(starts)
+        self.clip_v = clip_value
+        self.mode = mode
+
+        self.maps   = np.load(os.path.join(path,    'maps.npy'),    mmap_mode='c')
+        self.goals  = np.load(os.path.join(path,    'goals.npy'),   mmap_mode='c')
+        self.starts = np.load(os.path.join(path,    'starts.npy'),  mmap_mode='c')
 
-        if mode == 'f':
-            gt_values = np.load(os.path.join(path, 'focal.npy')).astype('float32')
-            gt_values = torch.tensor(gt_values)
-            self.gt_values = torch.where(gt_values >= clip_value, gt_values, torch.zeros_like(gt_values))
-        elif mode == 'h':
-            gt_values = np.load(os.path.join(path, 'abs.npy')).astype('float32')
-            self.gt_values = torch.tensor(gt_values)
-        elif mode == 'cf':
-            gt_values = np.load(os.path.join(path, 'cf.npy')).astype('float32')
-            self.gt_values = torch.tensor(gt_values)
-    
+        file_gt = {'f' : 'focal.npy', 'h':'abs.npy', 'cf': 'cf.npy'}[mode]
+        self.gt_values = np.load(os.path.join(path, file_gt), mmap_mode='c')
+
+
     def __len__(self):
         return len(self.gt_values)
 
+    
+    
     def __getitem__(self, idx):
-        return self.maps[idx], self.starts[idx], self.goals[idx], self.gt_values[idx]
+        gt_ = torch.from_numpy(self.gt_values[idx].astype('float32'))
+        if self.mode == 'f':
+            gt_=  torch.where( gt_ >= self.clip_v, gt_ , torch.zeros_like( torch.from_numpy(self.gt_values[idx])))
+        return (torch.from_numpy(self.maps[idx].astype('float32')), 
+                torch.from_numpy(self.starts[idx].astype('float32')), 
+                torch.from_numpy(self.goals[idx].astype('float32')), 
+                gt_ )

train.py

@@ -5,11 +5,13 @@
 import wandb
 from torch.utils.data import DataLoader
 from pytorch_lightning.loggers import WandbLogger
+import torch
 
 import argparse
+import multiprocessing
 
 
-def main(mode, run_name, proj_name):
+def main(mode, run_name, proj_name, batch_size, max_epochs):
     train_data = GridData(
         path='./TransPath_data/train',
         mode=mode
@@ -18,18 +20,25 @@ def main(mode, run_name, proj_name):
         path='./TransPath_data/val',
         mode=mode
     )
-    train_dataloader = DataLoader(train_data, batch_size=256,
-                        shuffle=True, num_workers=0, pin_memory=True)
-    val_dataloader = DataLoader(val_data, batch_size=256,
-                        shuffle=False, num_workers=0, pin_memory=True)
+    train_dataloader = DataLoader(  train_data, 
+                                    batch_size=batch_size,
+                                    shuffle=True, 
+                                    num_workers=multiprocessing.cpu_count(), 
+                                    pin_memory=True)
+    val_dataloader = DataLoader(    val_data, 
+                                    batch_size=batch_size,
+                                    shuffle=False, 
+                                    num_workers=multiprocessing.cpu_count(), 
+                                    pin_memory=True)
+    
     samples = next(iter(val_dataloader))
 
     model = Autoencoder(mode=mode)
-    wandb_logger = WandbLogger(project=proj_name, name=run_name)
+    wandb_logger = WandbLogger(project=proj_name, name=f'{run_name}_{mode}')
     trainer = pl.Trainer(
         logger=wandb_logger,
-        gpus=-1,
-        max_epochs=50,
+        accelerator="auto",
+        max_epochs=max_epochs,
         deterministic=False,
         callbacks=[PathLogger(samples, mode=mode)],
     )
@@ -42,12 +51,16 @@ def main(mode, run_name, proj_name):
     parser.add_argument('--run_name', type=str, default='default')
     parser.add_argument('--proj_name', type=str, default='TransPath_runs')
     parser.add_argument('--seed', type=int, default=39)
+    parser.add_argument('--batch', type=int, default=256)
+    parser.add_argument('--epoch', type=int, default=15)
 
     args = parser.parse_args()
     pl.seed_everything(args.seed)
-    
+    torch.set_float32_matmul_precision('high') #fix for tesor blocks warning with new video card
     main(
         mode=args.mode,
         run_name=args.run_name,
-        proj_name=args.proj_name
+        proj_name=args.proj_name,
+        batch_size=args.batch,
+        max_epochs=args.epoch
     )