Slight improvements to logging.

PiperOrigin-RevId: 359763463
google · Feb 26, 2021 · 889768b · 889768b
1 parent 78643ea
commit 889768b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/trax/rl/training.py b/trax/rl/training.py
@@ -259,7 +259,7 @@ def run(self, n_epochs=1, n_epochs_is_total_epochs=False):
             for steps in self._eval_steps:
               for eval_t in self._eval_temperatures:
                 sw.scalar(
-                    'rl/avg_return_temperature%d_steps%d' % (eval_t, steps),
+                    'rl/avg_return_temperature%.2f_steps%d' % (eval_t, steps),
                     self._avg_returns_temperatures[eval_t][steps][-1],
                     step=self._epoch)
           sw.scalar('rl/n_interactions', self.task.n_interactions(),

diff --git a/trax/supervised/training.py b/trax/supervised/training.py
@@ -402,6 +402,9 @@ def run(self, n_steps=1):
             (optimizer_metrics, loss))
 
         loss_acc += loss
+        # Log loss every 50 steps, every step in memory-efficient trainer.
+        if self._step % 50 == 0 or self._use_memory_efficient_trainer:
+          self._log_step('Loss: %.4f' % loss, stdout=False)
         step_acc += 1
         for metric_name, value in optimizer_metrics.items():
           optimizer_metrics_acc[metric_name] += value