diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index ee7c9e85..9d12ac81 100644 --- a/pytorch_lightning/core/lightning.py +++ b/pytorch_lightning/core/lightning.py @@ -92,7 +92,7 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): # ... return loss - # splitting it this way allows you to use your model as a feature extractor now + # splitting it this way allows model to be used a feature extractor model = MyModelAbove() inputs = server.get_request() @@ -565,11 +565,17 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): return model def init_ddp_connection(self, proc_rank, world_size): - """Connect all procs in the world using the env:// init - Use the first node as the root address + r""" - Override to init DDP in your own way. + Override to define your custom way of setting up a distributed environment. + Lightning's implementation uses env:// init by default and sets the first node as root. + + Args: + proc_rank (int): The current process rank within the node. + world_size (int): Number of GPUs being use across all nodes. (num_nodes*nb_gpu_nodes). + Example + ------- .. code-block:: python def init_ddp_connection(self): @@ -600,7 +606,11 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): root_node = self.trainer.resolve_root_node_address(root_node) os.environ['MASTER_ADDR'] = root_node - dist.init_process_group('nccl', rank=self.proc_rank, world_size=self.world_size) + dist.init_process_group( + 'nccl', + rank=self.proc_rank, + world_size=self.world_size + ) """ # use slurm job id for the port number @@ -945,29 +955,41 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): return [loader_a, loader_b, ..., loader_n] In the case where you return multiple `val_dataloaders`, the `validation_step` - will have an arguement `dataset_idx` which matches the order here. + will have an argument `dataset_idx` which matches the order here. """ return None @classmethod def load_from_metrics(cls, weights_path, tags_csv, map_location=None): - """Primary way of loading model from csv weights path. + r""" - :param str weights_path: Path to a PyTorch checkpoint - :param str tags_csv: Path to meta_tags.csv file generated by the test-tube Experiment - :param dict map_location: A dictionary mapping saved weight GPU devices to new GPU devices - for mapping storage {'cuda:1':'cuda:0'} - :return: The pretrained LightningModule + You should use `load_from_checkpoint` instead! + However, if your .ckpt weights don't have the hyperparameters saved, use this method to pass + in a .csv with the hparams you'd like to use. These will be converted into a argparse.Namespace + and passed into your LightningModule for use. - If you're using `test-tube`, there is an alternate method which uses the meta_tags.csv - file from test-tube to rebuild the model. The `meta_tags.csv` file can be found in the - `test-tube` experiment save_dir. + Args: + weights_path (str): Path to a PyTorch checkpoint + tags_csv (str): Path to a .csv with two columns (key, value) as in this + Example:: + key,value + drop_prob,0.2 + batch_size,32 + + map_location (dict): A dictionary mapping saved weight GPU devices to new + GPU devices (example: {'cuda:1':'cuda:0'}) + + Return: + LightningModule with loaded weights + + Example + ------- .. code-block:: python pretrained_model = MyLightningModule.load_from_metrics( weights_path='/path/to/pytorch_checkpoint.ckpt', - tags_csv='/path/to/test_tube/experiment/version/meta_tags.csv', + tags_csv='/path/to/hparams_file.csv', on_gpu=True, map_location=None ) @@ -976,22 +998,8 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): pretrained_model.eval() pretrained_model.freeze() y_hat = pretrained_model(x) - - This is the easiest/fastest way which loads hyperparameters and weights from a checkpoint, - such as the one saved by the `ModelCheckpoint` callback - - .. code-block:: python - - pretrained_model = MyLightningModule.load_from_checkpoint( - checkpoint_path='/path/to/pytorch_checkpoint.ckpt' - ) - - # predict - pretrained_model.eval() - pretrained_model.freeze() - y_hat = pretrained_model(x) - """ + hparams = load_hparams_from_tags_csv(tags_csv) hparams.__setattr__('on_gpu', False) @@ -1011,11 +1019,56 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): @classmethod def load_from_checkpoint(cls, checkpoint_path, map_location=None): - """ - Primary way of loading model from a checkpoint - :param checkpoint_path: - :param map_location: dic for mapping storage {'cuda:1':'cuda:0'} - :return: + r""" + + Primary way of loading model from a checkpoint. When Lightning saves a checkpoint + it stores the hyperparameters in the checkpoint if you initialized your LightningModule + with an argument called `hparams` which is a Namespace or dictionary of hyperparameters + + Example + ------- + .. code-block:: python + + # -------------- + # Case 1 + # when using Namespace (output of using Argparse to parse command line arguments) + from argparse import Namespace + hparams = Namespace(**{'learning_rate': 0.1}) + + model = MyModel(hparams) + + class MyModel(pl.LightningModule): + def __init__(self, hparams): + self.learning_rate = hparams.learning_rate + + # -------------- + # Case 2 + # when using a dict + model = MyModel({'learning_rate': 0.1}) + + class MyModel(pl.LightningModule): + def __init__(self, hparams): + self.learning_rate = hparams['learning_rate'] + + Args: + checkpoint_path (str): Path to checkpoint. + map_location (dic): If your checkpoint saved from a GPU model and you now load on CPUs + or a different number of GPUs, use this to map to the new setup. + + Return: + LightningModule with loaded weights. + + Example + ------- + .. code-block:: python + + # load weights without mapping + MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt') + + # load weights mapping all weights from GPU 1 to GPU 0 + map_location = {'cuda:1':'cuda:0'} + MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt', map_location=map_location) + """ if map_location is not None: @@ -1046,8 +1099,12 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): logging.info('\n' + model_summary.__str__()) def freeze(self): - """Freeze all params for inference + r""" + Freeze all params for inference + + Example + ------- .. code-block:: python model = MyLightningModule(...)