[RaySGD] Add tqdm logging to TorchTrainer (#7588)

* Update issue templates

* Init fp16

* fp16 and schedulers

* scheduler linking and fp16

* to fp16

* loss scaling and documentation

* more documentation

* add tests, refactor config

* moredocs

* more docs

* fix logo, add test mode, add fp16 flag

* fix tests

* fix scheduler

* fix apex

* improve safety

* fix tests

* fix tests

* remove pin memory default

* rm

* fix

* Update doc/examples/doc_code/raysgd_torch_signatures.py

* fix

* migrate changes from other PR

* ok thanks

* pass

* signatures

* lint'

* Update python/ray/experimental/sgd/pytorch/utils.py

* Apply suggestions from code review

Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com>

* should address most comments

* comments

* fix this ci

* first_pass

* add overrides

* override

* fixing up operators

* format

* sgd

* constants

* rm

* revert

* Checkpoint the basics

* End of day checkpoint

* Checkpoint log-to-head implementation

* Checkpoint

* Add actor-based batch log reporting, currently segfaults

* Work around progress segfault

* Fix some stuff in quicktorch

* Make things more customizable

* Quality of life fixes

* More quality of life

* Move tqdm logic to training_operator

* Update examples

* Fix some minor bugs

* Fix merge

* Fix small things, add pbar to dcgan

* Run format.sh

* Fix missing epoch number for batch pbar

* Address PR comments

* Fix float is not subscriptable

* Add train_loss to pbar by default

* Isolate tqdm code into a handler system

* Format

* Remove the batch_logs_reporter from distributed runner as well

* Check if the train_loss is avaialbale before using it

* Enable tqdm in the dcgan example

* Fix a crash in no-handler trainers

* Fix

* Allow not calling set_reporters for tests

Co-authored-by: Philipp Moritz <pcmoritz@gmail.com>
Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>

python/ray/util/sgd/torch/training_operator.py

@@ -1,4 +1,5 @@
 import collections
+
 import torch
 
 from ray.util.sgd.utils import (TimerCollection, AverageMeterCollection,
@@ -49,6 +50,9 @@ class TrainingOperator:
                  config,
                  models,
                  optimizers,
+                 train_loader,
+                 validation_loader,
+                 world_rank,
                  criterion=None,
                  schedulers=None,
                  use_fp16=False):
@@ -59,6 +63,9 @@ class TrainingOperator:
         self._optimizers = optimizers  # List of optimizers
         assert isinstance(optimizers, collections.Iterable), (
             "Components need to be iterable. Got: {}".format(type(optimizers)))
+        self._train_loader = train_loader
+        self._validation_loader = validation_loader
+        self._world_rank = world_rank
         self._criterion = criterion
         self._schedulers = schedulers
         if schedulers:
@@ -77,8 +84,12 @@ class TrainingOperator:
                         "TrainingOperator if using multi-scheduler, "
                         "multi-model or multi-optimizer training/validation.")
         self.timers = TimerCollection()
+        self.reporters = []
         self.setup(config)
 
+    def set_reporters(self, reporters):
+        self.reporters = reporters
+
     def _set_timers(self, timers):
         """Passes in the timers from the Runner."""
         self.timers = timers
@@ -131,6 +142,9 @@ class TrainingOperator:
         Returns:
             A dict of metrics from training.
         """
+        for r in self.reporters:
+            r.on_epoch_begin(info, self)
+
         metric_meters = AverageMeterCollection()
 
         self.model.train()
@@ -142,6 +156,9 @@ class TrainingOperator:
             batch_info.update(info)
             metrics = self.train_batch(batch, batch_info=batch_info)
 
+            for r in self.reporters:
+                r.on_batch_end(batch_info, metrics, self)
+
             if self.scheduler and batch_info.get(
                     SCHEDULER_STEP) == SCHEDULER_STEP_BATCH:
                 self.scheduler.step()
@@ -209,6 +226,7 @@ class TrainingOperator:
         # Call step of optimizer to update model params.
         with self.timers.record("apply"):
             self.optimizer.step()
+
         return {"train_loss": loss.item(), NUM_SAMPLES: features.size(0)}
 
     def validate(self, val_iterator, info):
@@ -318,6 +336,25 @@ class TrainingOperator:
         """List of optimizers created by the ``optimizer_creator``."""
         return self._optimizers
 
+    @property
+    def train_loader(self):
+        """
+        Data loader for the validation dataset created by the ``data_creator``.
+        """
+        return self._train_loader
+
+    @property
+    def validation_loader(self):
+        """
+        Data loader for the train dataset created by the ``data_creator``.
+        """
+        return self._validation_loader
+
+    @property
+    def world_rank(self):
+        """The rank of the parent runner. Always 0 if not distributed."""
+        return self._world_rank
+
     @property
     def criterion(self):
         """Criterion created by the provided ``loss_creator``."""