Custom argparser extension with Trainer arguments (argument types added)

CHANGELOG.md

@@ -25,7 +25,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
--
+- `Trainer.add_argparse_args` classmethod fixed. Now it adds a type for the arguments ([#1147](https://github.com/PyTorchLightning/pytorch-lightning/pull/1147)).
 
 ## [0.7.1] - 2020-03-07
 

pl_examples/full_examples/semantic_segmentation/models/unet/model.py

@@ -13,6 +13,7 @@ class UNet(nn.Module):
     bilinear (bool) - Whether to use bilinear interpolation or transposed
     convolutions for upsampling.
     '''
+
     def __init__(self, num_classes=19, bilinear=False):
         super().__init__()
         self.layer1 = DoubleConv(3, 64)

pl_examples/full_examples/semantic_segmentation/models/unet/parts.py

@@ -8,6 +8,7 @@ class DoubleConv(nn.Module):
     Double Convolution and BN and ReLU
     (3x3 conv -> BN -> ReLU) ** 2
     '''
+
     def __init__(self, in_ch, out_ch):
         super().__init__()
         self.net = nn.Sequential(
@@ -27,6 +28,7 @@ class Down(nn.Module):
     '''
     Combination of MaxPool2d and DoubleConv in series
     '''
+
     def __init__(self, in_ch, out_ch):
         super().__init__()
         self.net = nn.Sequential(
@@ -44,6 +46,7 @@ class Up(nn.Module):
     followed by concatenation of feature map from contracting path,
     followed by double 3x3 convolution.
     '''
+
     def __init__(self, in_ch, out_ch, bilinear=False):
         super().__init__()
         self.upsample = None

pl_examples/full_examples/semantic_segmentation/semseg.py

@@ -34,6 +34,7 @@ class KITTI(Dataset):
     encoded using `encode_segmap`, and given `transform` (if any) are applied to the image only
     (mask does not usually require transforms, but they can be implemented in a similar way).
     '''
+
     def __init__(
         self,
         root_path,
@@ -120,6 +121,7 @@ class SegModel(pl.LightningModule):
 
     Adam optimizer is used along with Cosine Annealing learning rate scheduler.
     '''
+
     def __init__(self, hparams):
         super(SegModel, self).__init__()
         self.root_path = hparams.root

pytorch_lightning/trainer/trainer.py

@@ -3,33 +3,33 @@
 import sys
 import warnings
 from argparse import ArgumentParser
-from typing import Union, Optional, List, Dict, Tuple, Iterable
+from typing import Union, Optional, List, Dict, Tuple, Iterable, Generator, Any
 
 import torch
-from torch import optim
 import torch.distributed as dist
 import torch.multiprocessing as mp
+from torch import optim
 from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 
 from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.loggers import LightningLoggerBase
 from pytorch_lightning.profiler import Profiler, PassThroughProfiler
 from pytorch_lightning.profiler.profiler import BaseProfiler
 from pytorch_lightning.trainer.auto_mix_precision import TrainerAMPMixin
 from pytorch_lightning.trainer.callback_config import TrainerCallbackConfigMixin
+from pytorch_lightning.trainer.callback_hook import TrainerCallbackHookMixin
 from pytorch_lightning.trainer.data_loading import TrainerDataLoadingMixin
+from pytorch_lightning.trainer.deprecated_api import TrainerDeprecatedAPITillVer0_8
 from pytorch_lightning.trainer.distrib_data_parallel import TrainerDDPMixin
 from pytorch_lightning.trainer.distrib_parts import (
     TrainerDPMixin,
     parse_gpu_ids,
     determine_root_gpu_device
 )
-from pytorch_lightning.core.lightning import LightningModule
-from pytorch_lightning.trainer.callback_hook import TrainerCallbackHookMixin
-from pytorch_lightning.trainer.deprecated_api import TrainerDeprecatedAPITillVer0_8
 from pytorch_lightning.trainer.evaluation_loop import TrainerEvaluationLoopMixin
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.trainer.model_hooks import TrainerModelHooksMixin
@@ -449,17 +449,52 @@ def default_attributes(cls):
         return args
 
     @classmethod
-    def add_argparse_args(cls, parent_parser: ArgumentParser) -> ArgumentParser:
-        """Extend existing argparse by default `Trainer` attributes."""
-        parser = ArgumentParser(parents=[parent_parser], add_help=False)
+    def _get_argparse_args_and_types(cls) -> Generator[Tuple[str, Any, Any], None, None]:
+        r"""Generates arguments, which could extend custom `ArgumentParser` object.
 
-        trainer_default_params = Trainer.default_attributes()
+        Yields:
+            Tuple with 3 values: argument name, argument type and argument default value.
+        """
+        import inspect
+        # Only arguments of these types can extend the `parent_parser`.
+        allowed_types = (str, float, int, bool)
+        trainer_default_params = inspect.signature(cls).parameters
 
-        # TODO: get "help" from docstring :)
         for arg in trainer_default_params:
+            arg_type = trainer_default_params[arg].annotation
+            arg_default = trainer_default_params[arg].default
+            try:
+                possible_arg_types = arg_type.__args__
+            except AttributeError:
+                possible_arg_types = (arg_type,)
+
+            allowed_arg_types = set(possible_arg_types).intersection(allowed_types)
+
+            arg_type = None
+            for allowed_type in allowed_types:
+                if allowed_type in allowed_arg_types:
+                    arg_type = allowed_type
+                    break
+
+            if arg_type is not None:
+                yield arg, arg_type, arg_default
+            else:
+                log.debug(
+                    'Argument %s has no allowed default type hint and will not be added '
+                    'to the arguments parser. Allowed types: %s', (arg, allowed_types)
+                )
+
+    @classmethod
+    def add_argparse_args(cls, parent_parser: ArgumentParser) -> ArgumentParser:
+        r"""Extend existing argparse by default `Trainer` attributes."""
+        parser = ArgumentParser(parents=[parent_parser], add_help=False, )
+
+        # TODO: get "help" from docstring :)
+        for arg, arg_type, arg_default in cls._get_argparse_args_and_types():
             parser.add_argument(
                 f'--{arg}',
-                default=trainer_default_params[arg],
+                default=arg_default,
+                type=arg_type,
                 dest=arg,
                 help='autogenerated by pl.Trainer'
             )
@@ -879,6 +914,7 @@ class _PatchDataLoader(object):
     Args:
         dataloader: Dataloader object to return when called.
     '''
+
     def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):
         self.dataloader = dataloader
 

tests/trainer/test_trainer.py

@@ -251,6 +251,7 @@ def test_dp_output_reduce():
 
 def test_model_checkpoint_options(tmpdir):
     """Test ModelCheckpoint options."""
+
     def mock_save_function(filepath):
         open(filepath, 'a').close()
 
@@ -644,3 +645,49 @@ def test_default_args(tmpdir):
 
     assert isinstance(trainer, Trainer)
     assert trainer.max_epochs == 5
+
+
+@pytest.mark.parametrize(
+    'cli_args',
+    [
+        ['--accumulate_grad_batches=22'],
+        ['--print_nan_grads=1', '--weights_save_path=./'],
+        []
+    ]
+)
+def test_add_argparse_args_redefined(cli_args):
+    """Redefines some default Trainer arguments via the cli and
+    tests the Trainer initialization correctness.
+    """
+    parser = ArgumentParser(add_help=False)
+    parser = Trainer.add_argparse_args(parent_parser=parser)
+
+    args = parser.parse_args(cli_args)
+
+    trainer = Trainer.from_argparse_args(args=args)
+    assert isinstance(trainer, Trainer)
+
+
+@pytest.mark.parametrize(
+    'cli_args',
+    [
+        ['--callbacks=1', '--logger'],
+        ['--foo', '--bar=1']
+    ]
+)
+def test_add_argparse_args_redefined_error(cli_args, monkeypatch):
+    """Asserts thar an error raised in case of passing not default cli arguments."""
+
+    class _UnkArgError(Exception):
+        pass
+
+    def _raise():
+        raise _UnkArgError
+
+    parser = ArgumentParser(add_help=False)
+    parser = Trainer.add_argparse_args(parent_parser=parser)
+
+    monkeypatch.setattr(parser, 'exit', lambda *args: _raise(), raising=True)
+
+    with pytest.raises(_UnkArgError):
+        parser.parse_args(cli_args)