review changes

Author: RaviTezu

ignite/contrib/engines/common.py

@@ -34,7 +34,7 @@
 def setup_common_training_handlers(
     trainer: Engine,
     train_sampler: Optional[DistributedSampler] = None,
-    to_save: Optional[Dict[str, Any]] = None,
+    to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
     lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
@@ -47,7 +47,7 @@ def setup_common_training_handlers(
     stop_on_nan: bool = True,
     clear_cuda_cache: bool = True,
     save_handler: Optional[Union[Callable, BaseSaveHandler]] = None,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Helper method to setup trainer with common handlers (it also supports distributed configuration):
         - :class:`~ignite.handlers.TerminateOnNan`
@@ -125,7 +125,7 @@ class to use to store ``to_save``. See :class:`~ignite.handlers.checkpoint.Check
 
 def _setup_common_training_handlers(
     trainer: Engine,
-    to_save: Optional[Dict[str, Any]] = None,
+    to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
     lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
@@ -137,7 +137,7 @@ def _setup_common_training_handlers(
     stop_on_nan: bool = True,
     clear_cuda_cache: bool = True,
     save_handler: Optional[Union[Callable, BaseSaveHandler]] = None,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     if output_path is not None and save_handler is not None:
         raise ValueError(
@@ -207,7 +207,7 @@ def output_transform(x, index, name):
 def _setup_common_distrib_training_handlers(
     trainer: Engine,
     train_sampler: Optional[DistributedSampler] = None,
-    to_save: Optional[Dict[str, Any]] = None,
+    to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
     lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
@@ -219,7 +219,7 @@ def _setup_common_distrib_training_handlers(
     stop_on_nan: bool = True,
     clear_cuda_cache: bool = True,
     save_handler: Optional[Union[Callable, BaseSaveHandler]] = None,
-    **kwargs: Any,
+    **kwargs: Any
 ):
 
     _setup_common_training_handlers(
@@ -313,7 +313,7 @@ def setup_tb_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup TensorBoard logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -344,7 +344,7 @@ def setup_visdom_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup Visdom logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -374,7 +374,7 @@ def setup_mlflow_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup MLflow logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -404,7 +404,7 @@ def setup_neptune_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup Neptune logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -434,7 +434,7 @@ def setup_wandb_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup WandB logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -464,7 +464,7 @@ def setup_plx_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup Polyaxon logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -494,7 +494,7 @@ def setup_trains_logging(
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     evaluators: Optional[Union[Engine, Dict[str, Engine]]] = None,
     log_every_iters: int = 100,
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method to setup Trains logging on trainer and a list of evaluators. Logged metrics are:
         - Training metrics, e.g. running average loss values
@@ -530,12 +530,12 @@ def wrapper(engine: Engine):
 def gen_save_best_models_by_val_score(
     save_handler: Union[Callable, BaseSaveHandler],
     evaluator: Engine,
-    models: torch.nn.Module,
+    models: Union[torch.nn.Module, Dict[str, torch.nn.Module]],
     metric_name: str,
     n_saved: int = 3,
     trainer: Optional[Engine] = None,
     tag: str = "val",
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method adds a handler to ``evaluator`` to save ``n_saved`` of best models based on the metric
     (named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``).
@@ -593,7 +593,7 @@ def save_best_model_by_val_score(
     n_saved: int = 3,
     trainer: Optional[Engine] = None,
     tag: str = "val",
-    **kwargs: Any,
+    **kwargs: Any
 ):
     """Method adds a handler to ``evaluator`` to save on a disk ``n_saved`` of best models based on the metric
     (named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``).

ignite/contrib/engines/tbptt.py

@@ -1,9 +1,8 @@
 # coding: utf-8
-from typing import Callable, Mapping, Optional, Sequence, Tuple, Union
+from typing import Callable, Mapping, Optional, Sequence, Union
 
 import torch
-from torch.nn import Module
-from torch.nn.modules.loss import _Loss
+import torch.nn as nn
 from torch.optim.optimizer import Optimizer
 
 from ignite.engine import Engine, EventEnum, _prepare_batch
@@ -31,9 +30,9 @@ def _detach_hidden(hidden: Union[torch.Tensor, Sequence, Mapping, str, bytes]):
 
 
 def create_supervised_tbptt_trainer(
-    model: Module,
+    model: nn.Module,
     optimizer: Optimizer,
-    loss_fn: _Loss,
+    loss_fn: nn.Module,
     tbtt_step: int,
     dim: int = 0,
     device: Optional[str] = None,

ignite/contrib/handlers/lr_finder.py

@@ -4,7 +4,7 @@
 import tempfile
 import warnings
 from pathlib import Path
-from typing import Any, Callable, Dict, Mapping, Optional
+from typing import Callable, Mapping, Optional
 
 import torch
 from torch.optim import Optimizer
@@ -254,7 +254,7 @@ def lr_suggestion(self):
     def attach(
         self,
         trainer: Engine,
-        to_save: Dict[str, Any],
+        to_save: Mapping,
         output_transform: Callable = lambda output: output,
         num_iter: Optional[int] = None,
         end_lr: float = 10.0,

ignite/contrib/handlers/param_scheduler.py

@@ -6,7 +6,7 @@
 from collections import OrderedDict
 from copy import copy
 from pathlib import Path
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
 from torch.optim.lr_scheduler import _LRScheduler
@@ -142,7 +142,7 @@ def get_param(self) -> Union[List[float], float]:
         pass
 
     @classmethod
-    def simulate_values(cls, num_events: int, **scheduler_kwargs: Mapping):
+    def simulate_values(cls, num_events: int, **scheduler_kwargs: Any):
         """Method to simulate scheduled values during `num_events` events.
 
         Args:
@@ -518,7 +518,7 @@ def state_dict(self):
             state_dict["schedulers"].append(s.state_dict())
         return state_dict
 
-    def load_state_dict(self, state_dict: Dict[str, Any]):
+    def load_state_dict(self, state_dict: Mapping):
         """Copies parameters from :attr:`state_dict` into this ConcatScheduler.
 
         Args:
@@ -583,7 +583,7 @@ def simulate_values(
         schedulers: List[ParamScheduler],
         durations: List[int],
         param_names: Union[List[str], Tuple[str]] = None,
-        **kwargs: Any,
+        **kwargs: Any
     ):
         """Method to simulate scheduled values during num_events events.
 
@@ -1048,7 +1048,7 @@ def state_dict(self):
             state_dict["schedulers"].append((n, s.state_dict()))
         return state_dict
 
-    def load_state_dict(self, state_dict: Dict):
+    def load_state_dict(self, state_dict: Mapping):
         """Copies parameters from :attr:`state_dict` into this ParamScheduler.
 
         Args:

ignite/contrib/handlers/time_profilers.py

@@ -388,7 +388,7 @@ def to_str(v: Union[str, tuple]):
                 return "{:.5f}/{}".format(v[0], v[1])
             return "{:.5f}".format(v)
 
-        def odict_to_str(d: Dict):
+        def odict_to_str(d: Mapping):
             out = " | ".join([to_str(v) for v in d.values()])
             return out
 

ignite/contrib/handlers/tqdm_logger.py

@@ -232,7 +232,7 @@ def __init__(
         description: str,
         metric_names: Optional[str] = None,
         output_transform: Optional[Callable] = None,
-        closing_event_name: Any = Events.EPOCH_COMPLETED,
+        closing_event_name: EventEnum = Events.EPOCH_COMPLETED,
     ):
         if metric_names is None and output_transform is None:
             # This helps to avoid 'Either metric_names or output_transform should be defined' of BaseOutputHandler

ignite/contrib/metrics/average_precision.py

@@ -24,7 +24,7 @@ class AveragePrecision(EpochMetric):
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
-        check_compute_fn (bool, optional): Optional default False. If True, `average_precision_score
+        check_compute_fn (bool): Default False. If True, `average_precision_score
             <http://scikit-learn.org/stable/modules/generated/sklearn.metrics.average_precision_score.html
             #sklearn.metrics.average_precision_score>`_ is run on the first batch of data to ensure there are
             no issues. User will be warned in case there are any issues computing the function.

ignite/contrib/metrics/precision_recall_curve.py

@@ -25,7 +25,7 @@ class PrecisionRecallCurve(EpochMetric):
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
-        check_compute_fn (bool, optional): Optional default False. If True, `precision_recall_curve
+        check_compute_fn (bool): Default False. If True, `precision_recall_curve
             <http://scikit-learn.org/stable/modules/generated/sklearn.metrics.precision_recall_curve.html
             #sklearn.metrics.precision_recall_curve>`_ is run on the first batch of data to ensure there are
             no issues. User will be warned in case there are any issues computing the function.

ignite/contrib/metrics/roc_auc.py

@@ -36,7 +36,7 @@ class ROC_AUC(EpochMetric):
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
-        check_compute_fn (bool, optional): Optional default False. If True, `roc_curve
+        check_compute_fn (bool): Default False. If True, `roc_curve
             <http://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_auc_score.html#
             sklearn.metrics.roc_auc_score>`_ is run on the first batch of data to ensure there are
             no issues. User will be warned in case there are any issues computing the function.
@@ -72,7 +72,7 @@ class RocCurve(EpochMetric):
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
             you want to compute the metric with respect to one of the outputs.
-        check_compute_fn (bool, optional): Optional default False. If True, `sklearn.metrics.roc_curve
+        check_compute_fn (bool): Default False. If True, `sklearn.metrics.roc_curve
             <http://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_curve.html#
             sklearn.metrics.roc_curve>`_ is run on the first batch of data to ensure there are
             no issues. User will be warned in case there are any issues computing the function.