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'