refactor: train_engine to trainer, eval_engine to evaluator (#76)

Author: ydcjeff

templates/_base/_handlers.py

@@ -16,8 +16,8 @@
 def get_handlers(
     config: Any,
     model: Module,
-    train_engine: Engine,
-    eval_engine: Engine,
+    trainer: Engine,
+    evaluator: Engine,
     metric_name: str,
     es_metric_name: str,
     train_sampler: Optional[DistributedSampler] = None,
@@ -48,9 +48,9 @@ def get_handlers(
 
     model
         best model to save
-    train_engine
+    trainer
         the engine used for training
-    eval_engine
+    evaluator
         the engine used for evaluation
     metric_name
         evaluation metric to save the best model
@@ -63,7 +63,7 @@ def get_handlers(
     lr_scheduler
         learning rate scheduler as native torch LRScheduler or ignite’s parameter scheduler
     output_names
-        list of names associated with `train_engine`'s process_function output dictionary
+        list of names associated with `trainer`'s process_function output dictionary
     kwargs
         keyword arguments passed to Checkpoint handler
 
@@ -78,7 +78,7 @@ def get_handlers(
     # kwargs can be passed to save the model based on training stats
     # like score_name, score_function
     common.setup_common_training_handlers(
-        trainer=train_engine,
+        trainer=trainer,
         train_sampler=train_sampler,
         to_save=to_save,
         lr_scheduler=lr_scheduler,
@@ -99,11 +99,11 @@ def get_handlers(
     # https://pytorch.org/ignite/contrib/engines.html#ignite.contrib.engines.common.save_best_model_by_val_score
     best_model_handler = common.save_best_model_by_val_score(
         output_path=config.output_dir / 'checkpoints',
-        evaluator=eval_engine,
+        evaluator=evaluator,
         model=model,
         metric_name=metric_name,
         n_saved=config.n_saved,
-        trainer=train_engine,
+        trainer=trainer,
         tag='eval',
     )
     {% endif %}
@@ -112,8 +112,8 @@ def get_handlers(
     # https://pytorch.org/ignite/contrib/engines.html#ignite.contrib.engines.common.add_early_stopping_by_val_score
     es_handler = common.add_early_stopping_by_val_score(
         patience=config.patience,
-        evaluator=eval_engine,
-        trainer=train_engine,
+        evaluator=evaluator,
+        trainer=trainer,
         metric_name=es_metric_name,
     )
     {% endif %}
@@ -125,7 +125,7 @@ def get_handlers(
     # you can replace with the events you want to measure
     timer_handler = Timer(average=True)
     timer_handler.attach(
-        engine=train_engine,
+        engine=trainer,
         start=Events.EPOCH_STARTED,
         resume=Events.ITERATION_STARTED,
         pause=Events.ITERATION_COMPLETED,
@@ -135,7 +135,7 @@ def get_handlers(
     {% if setup_timelimit %}
 
     # training will terminate if training time exceed `limit_sec`.
-    train_engine.add_event_handler(
+    trainer.add_event_handler(
         Events.ITERATION_COMPLETED, TimeLimit(limit_sec=config.limit_sec)
     )
     {% endif %}
@@ -144,8 +144,8 @@ def get_handlers(
 
 def get_logger(
     config: Any,
-    train_engine: Engine,
-    eval_engine: Optional[Union[Engine, Dict[str, Engine]]] = None,
+    trainer: Engine,
+    evaluator: Optional[Union[Engine, Dict[str, Engine]]] = None,
     optimizers: Optional[Union[Optimizer, Dict[str, Optimizer]]] = None,
     **kwargs: Any,
 ) -> Optional[BaseLogger]:
@@ -160,9 +160,9 @@ def get_logger(
     - `filepath`: logging path to output file
     - `logger_log_every_iters`: logging iteration interval for loggers
 
-    train_engine
+    trainer
         trainer engine
-    eval_engine
+    evaluator
         evaluator engine
     optimizers
         optimizers to log optimizer parameters
@@ -177,58 +177,58 @@ def get_logger(
 
     {% if logger_deps == 'clearml' %}
     logger_handler = common.setup_clearml_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'mlflow' %}
     logger_handler = common.setup_mlflow_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'neptune-client' %}
     logger_handler = common.setup_neptune_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'polyaxon-client' %}
     logger_handler = common.setup_plx_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'tensorboard' %}
     logger_handler = common.setup_tb_logging(
         output_path=config.output_dir,
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'visdom' %}
     logger_handler = common.setup_visdom_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )
     {% elif logger_deps == 'wandb' %}
     logger_handler = common.setup_wandb_logging(
-        trainer=train_engine,
+        trainer=trainer,
         optimizers=optimizers,
-        evaluators=eval_engine,
+        evaluators=evaluator,
         log_every_iters=config.logger_log_every_iters,
         **kwargs,
     )

templates/gan/_test_internal.py

@@ -27,7 +27,7 @@
 
 
 def test_get_handlers(tmp_path):
-    train_engine = Engine(lambda e, b: b)
+    trainer = Engine(lambda e, b: b)
     config = Namespace(
         output_dir=tmp_path,
         save_every_iters=1,
@@ -44,8 +44,8 @@ def test_get_handlers(tmp_path):
     bm_handler, es_handler, timer_handler = get_handlers(
         config=config,
         model=nn.Linear(1, 1),
-        train_engine=train_engine,
-        eval_engine=train_engine,
+        trainer=trainer,
+        evaluator=trainer,
         metric_name="eval_loss",
         es_metric_name="eval_loss",
     )
@@ -56,12 +56,12 @@ def test_get_handlers(tmp_path):
 
 def test_get_logger(tmp_path):
     config = Namespace(output_dir=tmp_path, logger_log_every_iters=1)
-    train_engine = Engine(lambda e, b: b)
+    trainer = Engine(lambda e, b: b)
     optimizer = optim.Adam(nn.Linear(1, 1).parameters())
     logger_handler = get_logger(
         config=config,
-        train_engine=train_engine,
-        eval_engine=train_engine,
+        trainer=trainer,
+        evaluator=trainer,
         optimizers=optimizer,
     )
     types = (
@@ -82,7 +82,7 @@ def test_create_trainers():
     model, optimizer, device, loss_fn, batch = set_up()
     real_labels = torch.ones(2, device=device)
     fake_labels = torch.zeros(2, device=device)
-    train_engine = create_trainers(
+    trainer = create_trainers(
         config=Namespace(use_amp=True),
         netD=model,
         netG=model,
@@ -93,7 +93,7 @@ def test_create_trainers():
         real_labels=real_labels,
         fake_labels=fake_labels,
     )
-    assert isinstance(train_engine, Engine)
+    assert isinstance(trainer, Engine)
 
 
 def test_get_default_parser():

templates/gan/main.py

@@ -71,14 +71,14 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
     netD, netG, optimizerD, optimizerG, loss_fn, lr_scheduler = initialize(config, num_channels)
 
     # -----------------------------
-    # train_engine and eval_engine
+    # trainer and evaluator
     # -----------------------------
     ws = idist.get_world_size()
     real_labels = torch.ones(config.batch_size // ws, device=device)
     fake_labels = torch.zeros(config.batch_size // ws, device=device)
     fixed_noise = torch.randn(config.batch_size // ws, config.z_dim, 1, 1, device=device)
 
-    train_engine = create_trainers(
+    trainer = create_trainers(
         config=config,
         netD=netD,
         netG=netG,
@@ -97,19 +97,19 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
 
     logger = setup_logging(config)
     log_basic_info(logger, config)
-    train_engine.logger = logger
+    trainer.logger = logger
 
     # -------------------------------------
     # ignite handlers and ignite loggers
     # -------------------------------------
 
-    to_save = {'netD': netD, 'netG': netG, 'optimizerD': optimizerD, 'optimizerG': optimizerG, 'trainer': train_engine}
+    to_save = {'netD': netD, 'netG': netG, 'optimizerD': optimizerD, 'optimizerG': optimizerG, 'trainer': trainer}
     optimizers = {'optimizerD': optimizerD, 'optimizerG': optimizerG}
     best_model_handler, es_handler, timer_handler = get_handlers(
         config=config,
         model={'netD', netD, 'netG', netG},
-        train_engine=train_engine,
-        eval_engine=train_engine,
+        trainer=trainer,
+        evaluator=trainer,
         metric_name='errD',
         es_metric_name='errD',
         to_save=to_save,
@@ -119,7 +119,7 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
 
     # setup ignite logger only on rank 0
     if rank == 0:
-        logger_handler = get_logger(config=config, train_engine=train_engine, optimizers=optimizers)
+        logger_handler = get_logger(config=config, trainer=trainer, optimizers=optimizers)
 
     # -----------------------------------
     # resume from the saved checkpoints
@@ -132,7 +132,7 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
     # adding handlers using `trainer.on` decorator API
     # --------------------------------------------------
 
-    @train_engine.on(Events.EPOCH_COMPLETED)
+    @trainer.on(Events.EPOCH_COMPLETED)
     def save_fake_example(engine):
         fake = netG(fixed_noise)
         path = config.output_dir / (FAKE_IMG_FNAME.format(engine.state.epoch))
@@ -141,7 +141,7 @@ def save_fake_example(engine):
     # --------------------------------------------------
     # adding handlers using `trainer.on` decorator API
     # --------------------------------------------------
-    @train_engine.on(Events.EPOCH_COMPLETED)
+    @trainer.on(Events.EPOCH_COMPLETED)
     def save_real_example(engine):
         img, y = engine.state.batch
         path = config.output_dir / (REAL_IMG_FNAME.format(engine.state.epoch))
@@ -150,13 +150,13 @@ def save_real_example(engine):
     # -------------------------------------------------------------
     # adding handlers using `trainer.on` decorator API
     # -------------------------------------------------------------
-    @train_engine.on(Events.EPOCH_COMPLETED)
+    @trainer.on(Events.EPOCH_COMPLETED)
     def print_times(engine):
         if not timer_handler:
             logger.info(f"Epoch {engine.state.epoch} done. Time per batch: {timer_handler.value():.3f}[s]")
             timer_handler.reset()
 
-    @train_engine.on(Events.ITERATION_COMPLETED(every=config.log_every_iters))
+    @trainer.on(Events.ITERATION_COMPLETED(every=config.log_every_iters))
     @idist.one_rank_only()
     def print_logs(engine):
         fname = config.output_dir / LOGS_FNAME
@@ -174,7 +174,7 @@ def print_logs(engine):
     # -------------------------------------------------------------
     # adding handlers using `trainer.on` decorator API
     # -------------------------------------------------------------
-    @train_engine.on(Events.EPOCH_COMPLETED)
+    @trainer.on(Events.EPOCH_COMPLETED)
     def create_plots(engine):
         try:
             import matplotlib as mpl
@@ -202,13 +202,13 @@ def create_plots(engine):
     # for training stats
     # --------------------------------
 
-    train_engine.add_event_handler(Events.ITERATION_COMPLETED(every=config.log_every_iters), log_metrics, tag="train")
+    trainer.add_event_handler(Events.ITERATION_COMPLETED(every=config.log_every_iters), log_metrics, tag="train")
 
     # ------------------------------------------
     # setup if done. let's run the training
     # ------------------------------------------
 
-    train_engine.run(train_dataloader, max_epochs=config.max_epochs, epoch_length=config.epoch_length)
+    trainer.run(train_dataloader, max_epochs=config.max_epochs, epoch_length=config.epoch_length)
 
     # ------------------------------------------------------------
     # close the logger after the training completed / terminated

templates/gan/trainers.py

@@ -1,5 +1,5 @@
 """
-`train_engine` and `eval_engine` like trainer and evaluator
+`trainer` and `evaluator` like trainer and evaluator
 """
 from typing import Any
 
@@ -131,13 +131,13 @@ def create_trainers(**kwargs) -> Engine:
 
     Returns
     -------
-    train_engine
+    trainer
     """
-    train_engine = Engine(
+    trainer = Engine(
         lambda e, b: train_function(
             engine=e,
             batch=b,
             **kwargs,
         )
     )
-    return train_engine
+    return trainer