added better documentation

lightly-ai · Jun 17, 2021 · c8e0a86 · c8e0a86
1 parent 498d643
commit c8e0a86
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/lightly/cli/config/config.yaml b/lightly/cli/config/config.yaml
@@ -71,7 +71,7 @@ trainer:
   gpus: 1                     # Number of gpus to use for training.
   max_epochs: 100             # Number of epochs to train for.
   precision: 32               # If set to 16, will use half-precision.
-  weights_summary: 'top'      # how to print the model architecture, one of {None, 'top', full},
+  weights_summary: 'top'      # how to print the model architecture, one of {None, top, full},
                                 #see https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#weights-summary
 
 # checkpoint_callback namespace: Modify the checkpoint callback

diff --git a/lightly/cli/lightly_cli.py b/lightly/cli/lightly_cli.py
@@ -22,7 +22,6 @@ def _lightly_cli(cfg, is_cli_call=True):
     if cfg['trainer']['max_epochs'] > 0:
         print('#' * 10 + ' Starting to train an embedding model.')
         checkpoint = _train_cli(cfg, is_cli_call)
-        cfg['trainer']['weights_summary'] = None
     else:
         checkpoint = ''
 

diff --git a/lightly/cli/train_cli.py b/lightly/cli/train_cli.py
@@ -157,6 +157,9 @@ def train_cli(cfg):
         >>>
         >>> # train model for 10 epochs
         >>> lightly-train input_dir=data/ trainer.max_epochs=10
+        >>>
+        >>> # print a full summary of the model
+        >>> lightly-train input_dir=data/ trainer.weights_summary=full
 
     """
     return _train_cli(cfg)