[tune] Add PyTorch MNIST Example + Misc. Tweaks (#2708)

python/ray/tune/trainable.py

@@ -74,7 +74,7 @@ class Trainable(object):
 
         self._iteration = 0
         self._time_total = 0.0
-        self._timesteps_total = 0
+        self._timesteps_total = None
         self._setup()
         self._initialize_ok = True
         self._local_ip = ray.services.get_node_ip_address()
@@ -150,9 +150,15 @@ class Trainable(object):
             time_this_iter = time.time() - start
         self._time_total += time_this_iter
 
-        self._timesteps_total += result.get(TIMESTEPS_THIS_ITER, 0)
-
         result.setdefault(DONE, False)
+
+        # self._timesteps_total should only be tracked if increments provided
+        if result.get(TIMESTEPS_THIS_ITER):
+            if self._timesteps_total is None:
+                self._timesteps_total = 0
+            self._timesteps_total += result[TIMESTEPS_THIS_ITER]
+
+        # self._timesteps_total should not override user-provided total
         result.setdefault(TIMESTEPS_TOTAL, self._timesteps_total)
 
         # Provides auto-filled neg_mean_loss for avoiding regressions
@@ -278,12 +284,26 @@ class Trainable(object):
         raise NotImplementedError
 
     def _save(self, checkpoint_dir):
-        """Subclasses should override this to implement save()."""
+        """Subclasses should override this to implement save().
+
+        Args:
+            checkpoint_dir (str): The directory where the checkpoint
+                can be stored.
+
+        Returns:
+            Checkpoint path that may be passed to restore(). Typically
+                would default to `checkpoint_dir`.
+        """
 
         raise NotImplementedError
 
     def _restore(self, checkpoint_path):
-        """Subclasses should override this to implement restore()."""
+        """Subclasses should override this to implement restore().
+
+        Args:
+            checkpoint_path (str): The directory where the checkpoint
+                is stored.
+        """
 
         raise NotImplementedError