Metrics load (#228)

* load from metrics defaults to CPU

* load from metrics defaults to CPU

* load from metrics defaults to CPU

Author: williamFalcon

pytorch_lightning/root_module/root_module.py

@@ -130,7 +130,7 @@ def val_dataloader(self):
         return None
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
@@ -142,13 +142,9 @@ def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
-        if on_gpu:
-            if map_location is not None:
-                checkpoint = torch.load(weights_path, map_location=map_location)
-            else:
-                checkpoint = torch.load(weights_path)
-        else:
-            checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
+        # load on CPU only to avoid OOM issues
+        # then its up to user to put back on GPUs
+        checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
         # load the state_dict on the model automatically
         model = cls(hparams)

tests/test_models.py

@@ -1030,8 +1030,7 @@ def test_amp_gpu_ddp_slurm_managed():
     assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
 
     # test model loading with a map_location
-    map_location = 'cuda:1'
-    pretrained_model = load_model(exp, save_dir, True, map_location)
+    pretrained_model = load_model(exp, save_dir, True)
 
     # test model preds
     run_prediction(model.test_dataloader, pretrained_model)
@@ -1406,7 +1405,7 @@ def clear_save_dir():
         shutil.move(save_dir, save_dir + f'_{n}')
 
 
-def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel):
+def load_model(exp, save_dir, on_gpu, module_class=LightningTemplateModel):
 
     # load trained model
     tags_path = exp.get_data_path(exp.name, exp.version)
@@ -1417,8 +1416,7 @@ def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningT
 
     trained_model = module_class.load_from_metrics(weights_path=weights_dir,
                                                    tags_csv=tags_path,
-                                                   on_gpu=on_gpu,
-                                                   map_location=map_location)
+                                                   on_gpu=on_gpu)
 
     assert trained_model is not None, 'loading model failed'