cleared spaces

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(...)