From 52db52b927d8050f25eeaec162ac715c42a6da9a Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Wed, 3 Jan 2024 09:06:26 +0200
Subject: [PATCH] Add ndjson logging for training (#10970)

* Add ndjson logging for training

This adds support for NDJSON (newline-delimited JSON) metrics logging,
for both console (stdout) output and a file (like the current CSV file).

NDJSON can be easily grepped from the output and/or parsed with e.g. `jq`.

The feature is enabled with the `--ndjson-console` and `--ndjson-file`
switches to `train.py`.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 train.py                  | 21 +++++++++++++++++++--
 utils/loggers/__init__.py | 24 +++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/train.py b/train.py
index 4c3bec3..ad37f2c 100644
--- a/train.py
+++ b/train.py
@@ -58,7 +58,7 @@ from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, ch
                            get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
                            labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,
                            yaml_save)
-from utils.loggers import Loggers
+from utils.loggers import LOGGERS, Loggers
 from utils.loggers.comet.comet_utils import check_comet_resume
 from utils.loss import ComputeLoss
 from utils.metrics import fitness
@@ -98,7 +98,20 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     # Loggers
     data_dict = None
     if RANK in {-1, 0}:
-        loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)  # loggers instance
+        include_loggers = list(LOGGERS)
+        if getattr(opt, 'ndjson_console', False):
+            include_loggers.append('ndjson_console')
+        if getattr(opt, 'ndjson_file', False):
+            include_loggers.append('ndjson_file')
+
+        loggers = Loggers(
+            save_dir=save_dir,
+            weights=weights,
+            opt=opt,
+            hyp=hyp,
+            logger=LOGGER,
+            include=tuple(include_loggers),
+        )
 
         # Register actions
         for k in methods(loggers):
@@ -482,6 +495,10 @@ def parse_opt(known=False):
     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
     parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
 
+    # NDJSON logging
+    parser.add_argument('--ndjson-console', action='store_true', help='Log ndjson to console')
+    parser.add_argument('--ndjson-file', action='store_true', help='Log ndjson to file')
+
     return parser.parse_known_args()[0] if known else parser.parse_args()
 
 
diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
index ba7d279..aa6eae4 100644
--- a/utils/loggers/__init__.py
+++ b/utils/loggers/__init__.py
@@ -2,7 +2,7 @@
 """
 Logging utils
 """
-
+import json
 import os
 import warnings
 from pathlib import Path
@@ -58,6 +58,18 @@ except (ImportError, AssertionError):
     comet_ml = None
 
 
+def _json_default(value):
+    """Format `value` for JSON serialization (e.g. unwrap tensors). Fall back to strings."""
+    if isinstance(value, torch.Tensor):
+        try:
+            value = value.item()
+        except ValueError:  # "only one element tensors can be converted to Python scalars"
+            pass
+    if isinstance(value, float):
+        return value
+    return str(value)
+
+
 class Loggers():
     # YOLOv5 Loggers class
     def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
@@ -86,6 +98,8 @@ class Loggers():
         for k in LOGGERS:
             setattr(self, k, None)  # init empty logger dictionary
         self.csv = True  # always log to csv
+        self.ndjson_console = ('ndjson_console' in self.include)  # log ndjson to console
+        self.ndjson_file = ('ndjson_file' in self.include)  # log ndjson to file
 
         # Messages
         if not comet_ml:
@@ -228,6 +242,14 @@ class Loggers():
             s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n')  # add header
             with open(file, 'a') as f:
                 f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
+        if self.ndjson_console or self.ndjson_file:
+            json_data = json.dumps(dict(epoch=epoch, **x), default=_json_default)
+            if self.ndjson_console:
+                print(json_data)
+            if self.ndjson_file:
+                file = self.save_dir / 'results.ndjson'
+                with open(file, 'a') as f:
+                    print(json_data, file=f)
 
         if self.tb:
             for k, v in x.items():