[trainer] add Trainer methods for metrics logging and saving (huggingface#10266)

stas00 · sgugger · web-flow · commit 622a8c5995ba · 2021-02-22T13:02:53.000-08:00
* make logging and saving trainer built-in

* Update src/transformers/trainer.py

Co-authored-by: Sylvain Gugger &lt;35901082+sgugger@users.noreply.github.com&gt;

Co-authored-by: Sylvain Gugger &lt;35901082+sgugger@users.noreply.github.com&gt;
diff --git a/examples/seq2seq/run_seq2seq.py b/examples/seq2seq/run_seq2seq.py
@@ -18,7 +18,6 @@
 """
 # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
 
-import json
 import logging
 import os
 import re
@@ -55,11 +54,6 @@
 logger = logging.getLogger(__name__)
 
 
-def save_json(content, path, indent=4, **json_dump_kwargs):
-    with open(path, "w") as f:
-        json.dump(content, f, indent=indent, sort_keys=True, **json_dump_kwargs)
-
-
 @dataclass
 class ModelArguments:
     """
@@ -596,13 +590,8 @@ def compute_metrics(eval_preds):
         )
         metrics["train_samples"] = min(max_train_samples, len(train_dataset))
         if trainer.is_world_process_zero():
-            metrics_formatted = trainer.metrics_format(metrics)
-            logger.info("***** train metrics *****")
-            k_width = max(len(str(x)) for x in metrics_formatted.keys())
-            v_width = max(len(str(x)) for x in metrics_formatted.values())
-            for key in sorted(metrics_formatted.keys()):
-                logger.info(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")
-            save_json(metrics, os.path.join(training_args.output_dir, "train_results.json"))
+            trainer.log_metrics("train", metrics)
+            trainer.save_metrics("train", metrics)
             all_metrics.update(metrics)
 
             # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
@@ -620,13 +609,8 @@ def compute_metrics(eval_preds):
         metrics["eval_samples"] = min(max_val_samples, len(eval_dataset))
 
         if trainer.is_world_process_zero():
-            metrics_formatted = trainer.metrics_format(metrics)
-            logger.info("***** val metrics *****")
-            k_width = max(len(str(x)) for x in metrics_formatted.keys())
-            v_width = max(len(str(x)) for x in metrics_formatted.values())
-            for key in sorted(metrics_formatted.keys()):
-                logger.info(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")
-            save_json(metrics, os.path.join(training_args.output_dir, "eval_results.json"))
+            trainer.log_metrics("eval", metrics)
+            trainer.save_metrics("eval", metrics)
             all_metrics.update(metrics)
 
     if training_args.do_predict:
@@ -643,13 +627,8 @@ def compute_metrics(eval_preds):
         metrics["test_samples"] = min(max_test_samples, len(test_dataset))
 
         if trainer.is_world_process_zero():
-            metrics_formatted = trainer.metrics_format(metrics)
-            logger.info("***** test metrics *****")
-            k_width = max(len(str(x)) for x in metrics_formatted.keys())
-            v_width = max(len(str(x)) for x in metrics_formatted.values())
-            for key in sorted(metrics_formatted.keys()):
-                logger.info(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")
-            save_json(metrics, os.path.join(training_args.output_dir, "test_results.json"))
+            trainer.log_metrics("test", metrics)
+            trainer.save_metrics("test", metrics)
             all_metrics.update(metrics)
 
             if training_args.predict_with_generate:
@@ -662,7 +641,7 @@ def compute_metrics(eval_preds):
                     writer.write("\n".join(test_preds))
 
     if trainer.is_world_process_zero():
-        save_json(all_metrics, os.path.join(training_args.output_dir, "all_results.json"))
+        trainer.save_metrics("all", metrics)
 
     return results
 
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -19,6 +19,7 @@
 import collections
 import gc
 import inspect
+import json
 import math
 import os
 import re
@@ -1370,6 +1371,38 @@ def metrics_format(self, metrics: Dict[str, float]) -> Dict[str, float]:
 
         return metrics_copy
 
+    def log_metrics(self, split, metrics):
+        """
+        Log metrics in a specially formatted way
+
+        Args:
+            split (:obj:`str`):
+                Mode/split name: one of ``train``, ``eval``, ``test``
+            metrics (:obj:`Dict[str, float]`):
+                The metrics returned from train/evaluate/predictmetrics: metrics dict
+        """
+
+        logger.info(f"***** {split} metrics *****")
+        metrics_formatted = self.metrics_format(metrics)
+        k_width = max(len(str(x)) for x in metrics_formatted.keys())
+        v_width = max(len(str(x)) for x in metrics_formatted.values())
+        for key in sorted(metrics_formatted.keys()):
+            logger.info(f"  {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}")
+
+    def save_metrics(self, split, metrics):
+        """
+        Save metrics into a json file for that split, e.g. ``train_results.json``.
+
+        Args:
+            split (:obj:`str`):
+                Mode/split name: one of ``train``, ``eval``, ``test``, ``all``
+            metrics (:obj:`Dict[str, float]`):
+                The metrics returned from train/evaluate/predict
+        """
+        path = os.path.join(self.args.output_dir, f"{split}_results.json")
+        with open(path, "w") as f:
+            json.dump(metrics, f, indent=4, sort_keys=True)
+
     def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
         """
         Prepare :obj:`inputs` before feeding them to the model, converting them to tensors if they are not already and
