chore(logger): log predictions during training to wandb tables (#1181)

README.md

@@ -150,6 +150,19 @@ On the second machine, run
 python tools/train.py -n yolox-s -b 128 --dist-url tcp://123.123.123.123:12312 --num_machines 2 --machine_rank 1
 ```
 
+**Logging to Weights & Biases**
+
+To log metrics, predictions and model checkpoints to [W&B](https://docs.wandb.ai/guides/integrations/other/yolox) use the command line argument `--logger wandb` and use the prefix "wandb-" to specify arguments for initializing the wandb run.
+
+```shell
+python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb wandb-project <project name>
+                         yolox-m
+                         yolox-l
+                         yolox-x
+```
+
+An example wandb dashboard is available [here](https://wandb.ai/manan-goel/yolox-nano/runs/3pzfeom0)
+
 **Others**  
 See more information with the following command:
 ```shell

docs/quick_run.md

@@ -76,6 +76,19 @@ python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb w
                          yolox-x
 ```
 
+More WandbLogger arguments include
+
+```shell
+python tools/train.py .... --logger wandb wandb-project <project-name> \
+                wandb-name <run-name> \
+                wandb-id <run-id> \
+                wandb-save_dir <save-dir> \
+                wandb-num_eval_images <num-images> \
+                wandb-log_checkpoints <bool>
+```
+
+More information available [here](https://docs.wandb.ai/guides/integrations/other/yolox).
+
 **Multi Machine Training**
 
 We also support multi-nodes training. Just add the following args:

tools/train.py

@@ -84,7 +84,8 @@ def make_parser():
         "-l",
         "--logger",
         type=str,
-        help="Logger to be used for metrics",
+        help="Logger to be used for metrics. \
+        Implemented loggers include `tensorboard` and `wandb`.",
         default="tensorboard"
     )
     parser.add_argument(

yolox/core/trainer.py

@@ -180,11 +180,11 @@ class Trainer:
             if self.args.logger == "tensorboard":
                 self.tblogger = SummaryWriter(os.path.join(self.file_name, "tensorboard"))
             elif self.args.logger == "wandb":
-                wandb_params = dict()
-                for k, v in zip(self.args.opts[0::2], self.args.opts[1::2]):
-                    if k.startswith("wandb-"):
-                        wandb_params.update({k[len("wandb-"):]: v})
-                self.wandb_logger = WandbLogger(config=vars(self.exp), **wandb_params)
+                self.wandb_logger = WandbLogger.initialize_wandb_logger(
+                    self.args,
+                    self.exp,
+                    self.evaluator.dataloader.dataset
+                )
             else:
                 raise ValueError("logger must be either 'tensorboard' or 'wandb'")
 
@@ -263,8 +263,11 @@ class Trainer:
 
             if self.rank == 0:
                 if self.args.logger == "wandb":
-                    self.wandb_logger.log_metrics({k: v.latest for k, v in loss_meter.items()})
-                    self.wandb_logger.log_metrics({"lr": self.meter["lr"].latest})
+                    metrics = {"train/" + k: v.latest for k, v in loss_meter.items()}
+                    metrics.update({
+                        "train/lr": self.meter["lr"].latest
+                    })
+                    self.wandb_logger.log_metrics(metrics, step=self.progress_in_iter)
 
             self.meter.clear_meters()
 
@@ -322,8 +325,8 @@ class Trainer:
                 evalmodel = evalmodel.module
 
         with adjust_status(evalmodel, training=False):
-            ap50_95, ap50, summary = self.exp.eval(
-                evalmodel, self.evaluator, self.is_distributed
+            (ap50_95, ap50, summary), predictions = self.exp.eval(
+                evalmodel, self.evaluator, self.is_distributed, return_outputs=True
             )
 
         update_best_ckpt = ap50_95 > self.best_ap
@@ -337,16 +340,17 @@ class Trainer:
                 self.wandb_logger.log_metrics({
                     "val/COCOAP50": ap50,
                     "val/COCOAP50_95": ap50_95,
-                    "epoch": self.epoch + 1,
+                    "train/epoch": self.epoch + 1,
                 })
+                self.wandb_logger.log_images(predictions)
             logger.info("\n" + summary)
         synchronize()
 
-        self.save_ckpt("last_epoch", update_best_ckpt)
+        self.save_ckpt("last_epoch", update_best_ckpt, ap=ap50_95)
         if self.save_history_ckpt:
-            self.save_ckpt(f"epoch_{self.epoch + 1}")
+            self.save_ckpt(f"epoch_{self.epoch + 1}", ap=ap50_95)
 
-    def save_ckpt(self, ckpt_name, update_best_ckpt=False):
+    def save_ckpt(self, ckpt_name, update_best_ckpt=False, ap=None):
         if self.rank == 0:
             save_model = self.ema_model.ema if self.use_model_ema else self.model
             logger.info("Save weights to {}".format(self.file_name))
@@ -355,6 +359,7 @@ class Trainer:
                 "model": save_model.state_dict(),
                 "optimizer": self.optimizer.state_dict(),
                 "best_ap": self.best_ap,
+                "curr_ap": ap,
             }
             save_checkpoint(
                 ckpt_state,
@@ -364,4 +369,14 @@ class Trainer:
             )
 
             if self.args.logger == "wandb":
-                self.wandb_logger.save_checkpoint(self.file_name, ckpt_name, update_best_ckpt)
+                self.wandb_logger.save_checkpoint(
+                    self.file_name,
+                    ckpt_name,
+                    update_best_ckpt,
+                    metadata={
+                        "epoch": self.epoch + 1,
+                        "optimizer": self.optimizer.state_dict(),
+                        "best_ap": self.best_ap,
+                        "curr_ap": ap
+                    }
+                )

yolox/data/datasets/coco.py

@@ -65,8 +65,8 @@ class COCODataset(Dataset):
         remove_useless_info(self.coco)
         self.ids = self.coco.getImgIds()
         self.class_ids = sorted(self.coco.getCatIds())
-        cats = self.coco.loadCats(self.coco.getCatIds())
-        self._classes = tuple([c["name"] for c in cats])
+        self.cats = self.coco.loadCats(self.coco.getCatIds())
+        self._classes = tuple([c["name"] for c in self.cats])
         self.imgs = None
         self.name = name
         self.img_size = img_size

yolox/evaluators/coco_evaluator.py

@@ -8,6 +8,7 @@ import itertools
 import json
 import tempfile
 import time
+from collections import ChainMap, defaultdict
 from loguru import logger
 from tabulate import tabulate
 from tqdm import tqdm
@@ -120,6 +121,7 @@ class COCOEvaluator:
         trt_file=None,
         decoder=None,
         test_size=None,
+        return_outputs=False
     ):
         """
         COCO average precision (AP) Evaluation. Iterate inference on the test dataset
@@ -142,6 +144,7 @@ class COCOEvaluator:
             model = model.half()
         ids = []
         data_list = []
+        output_data = defaultdict()
         progress_bar = tqdm if is_main_process() else iter
 
         inference_time = 0
@@ -184,20 +187,29 @@ class COCOEvaluator:
                     nms_end = time_synchronized()
                     nms_time += nms_end - infer_end
 
-            data_list.extend(self.convert_to_coco_format(outputs, info_imgs, ids))
+            data_list_elem, image_wise_data = self.convert_to_coco_format(
+                outputs, info_imgs, ids, return_outputs=True)
+            data_list.extend(data_list_elem)
+            output_data.update(image_wise_data)
 
         statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])
         if distributed:
             data_list = gather(data_list, dst=0)
+            output_data = gather(output_data, dst=0)
             data_list = list(itertools.chain(*data_list))
+            output_data = dict(ChainMap(*output_data))
             torch.distributed.reduce(statistics, dst=0)
 
         eval_results = self.evaluate_prediction(data_list, statistics)
         synchronize()
+
+        if return_outputs:
+            return eval_results, output_data
         return eval_results
 
-    def convert_to_coco_format(self, outputs, info_imgs, ids):
+    def convert_to_coco_format(self, outputs, info_imgs, ids, return_outputs=False):
         data_list = []
+        image_wise_data = defaultdict(dict)
         for (output, img_h, img_w, img_id) in zip(
             outputs, info_imgs[0], info_imgs[1], ids
         ):
@@ -212,10 +224,22 @@ class COCOEvaluator:
                 self.img_size[0] / float(img_h), self.img_size[1] / float(img_w)
             )
             bboxes /= scale
-            bboxes = xyxy2xywh(bboxes)
-
             cls = output[:, 6]
             scores = output[:, 4] * output[:, 5]
+
+            image_wise_data.update({
+                int(img_id): {
+                    "bboxes": [box.numpy().tolist() for box in bboxes],
+                    "scores": [score.numpy().item() for score in scores],
+                    "categories": [
+                        self.dataloader.dataset.class_ids[int(cls[ind])]
+                        for ind in range(bboxes.shape[0])
+                    ],
+                }
+            })
+
+            bboxes = xyxy2xywh(bboxes)
+
             for ind in range(bboxes.shape[0]):
                 label = self.dataloader.dataset.class_ids[int(cls[ind])]
                 pred_data = {
@@ -226,6 +250,9 @@ class COCOEvaluator:
                     "segmentation": [],
                 }  # COCO json format
                 data_list.append(pred_data)
+
+        if return_outputs:
+            return data_list, image_wise_data
         return data_list
 
     def evaluate_prediction(self, data_dict, statistics):

yolox/exp/yolox_base.py

@@ -318,5 +318,5 @@ class Exp(BaseExp):
         # NOTE: trainer shouldn't be an attribute of exp object
         return trainer
 
-    def eval(self, model, evaluator, is_distributed, half=False):
-        return evaluator.evaluate(model, is_distributed, half)
+    def eval(self, model, evaluator, is_distributed, half=False, return_outputs=False):
+        return evaluator.evaluate(model, is_distributed, half, return_outputs=return_outputs)

yolox/utils/logger.py

@@ -5,8 +5,12 @@
 import inspect
 import os
 import sys
+from collections import defaultdict
 from loguru import logger
 
+import cv2
+import numpy as np
+
 import torch
 
 
@@ -108,6 +112,7 @@ class WandbLogger(object):
 
     For more information, please refer to:
     https://docs.wandb.ai/guides/track
+    https://docs.wandb.ai/guides/integrations/other/yolox
     """
     def __init__(self,
                  project=None,
@@ -116,6 +121,9 @@ class WandbLogger(object):
                  entity=None,
                  save_dir=None,
                  config=None,
+                 val_dataset=None,
+                 num_eval_images=100,
+                 log_checkpoints=False,
                  **kwargs):
         """
         Args:
@@ -125,7 +133,24 @@ class WandbLogger(object):
             entity (str): wandb entity name.
             save_dir (str): save directory.
             config (dict): config dict.
+            val_dataset (Dataset): validation dataset.
+            num_eval_images (int): number of images from the validation set to log.
+            log_checkpoints (bool): log checkpoints
             **kwargs: other kwargs.
+
+        Usage:
+            Any arguments for wandb.init can be provided on the command line using
+            the prefix `wandb-`.
+            Example
+            ```
+            python tools/train.py .... --logger wandb wandb-project <project-name> \
+                wandb-name <run-name> \
+                wandb-id <run-id> \
+                wandb-save_dir <save-dir> \
+                wandb-num_eval_imges <num-images> \
+                wandb-log_checkpoints <bool>
+            ```
+            The val_dataset argument is not open to the command line.
         """
         try:
             import wandb
@@ -144,6 +169,12 @@ class WandbLogger(object):
         self.kwargs = kwargs
         self.entity = entity
         self._run = None
+        self.val_artifact = None
+        if num_eval_images == -1:
+            self.num_log_images = len(val_dataset)
+        else:
+            self.num_log_images = min(num_eval_images, len(val_dataset))
+        self.log_checkpoints = (log_checkpoints == "True" or log_checkpoints == "true")
         self._wandb_init = dict(
             project=self.project,
             name=self.name,
@@ -158,8 +189,17 @@ class WandbLogger(object):
 
         if self.config:
             self.run.config.update(self.config)
-        self.run.define_metric("epoch")
-        self.run.define_metric("val/", step_metric="epoch")
+        self.run.define_metric("train/epoch")
+        self.run.define_metric("val/*", step_metric="train/epoch")
+        self.run.define_metric("train/step")
+        self.run.define_metric("train/*", step_metric="train/step")
+
+        if val_dataset and self.num_log_images != 0:
+            self.cats = val_dataset.cats
+            self.id_to_class = {
+                cls['id']: cls['name'] for cls in self.cats
+            }
+            self._log_validation_set(val_dataset)
 
     @property
     def run(self):
@@ -176,6 +216,32 @@ class WandbLogger(object):
                 self._run = self.wandb.init(**self._wandb_init)
         return self._run
 
+    def _log_validation_set(self, val_dataset):
+        """
+        Log validation set to wandb.
+
+        Args:
+            val_dataset (Dataset): validation dataset.
+        """
+        if self.val_artifact is None:
+            self.val_artifact = self.wandb.Artifact(name="validation_images", type="dataset")
+            self.val_table = self.wandb.Table(columns=["id", "input"])
+
+            for i in range(self.num_log_images):
+                data_point = val_dataset[i]
+                img = data_point[0]
+                id = data_point[3]
+                img = np.transpose(img, (1, 2, 0))
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                self.val_table.add_data(
+                    id.item(),
+                    self.wandb.Image(img)
+                )
+
+            self.val_artifact.add(self.val_table, "validation_images_table")
+            self.run.use_artifact(self.val_artifact)
+            self.val_artifact.wait()
+
     def log_metrics(self, metrics, step=None):
         """
         Args:
@@ -188,21 +254,98 @@ class WandbLogger(object):
                 metrics[k] = v.item()
 
         if step is not None:
-            self.run.log(metrics, step=step)
+            metrics.update({"train/step": step})
+            self.run.log(metrics)
         else:
             self.run.log(metrics)
 
-    def save_checkpoint(self, save_dir, model_name, is_best):
+    def log_images(self, predictions):
+        if len(predictions) == 0 or self.val_artifact is None or self.num_log_images == 0:
+            return
+
+        table_ref = self.val_artifact.get("validation_images_table")
+
+        columns = ["id", "predicted"]
+        for cls in self.cats:
+            columns.append(cls["name"])
+
+        result_table = self.wandb.Table(columns=columns)
+        for idx, val in table_ref.iterrows():
+
+            avg_scores = defaultdict(int)
+            num_occurrences = defaultdict(int)
+
+            if val[0] in predictions:
+                prediction = predictions[val[0]]
+                boxes = []
+
+                for i in range(len(prediction["bboxes"])):
+                    bbox = prediction["bboxes"][i]
+                    x0 = bbox[0]
+                    y0 = bbox[1]
+                    x1 = bbox[2]
+                    y1 = bbox[3]
+                    box = {
+                        "position": {
+                            "minX": min(x0, x1),
+                            "minY": min(y0, y1),
+                            "maxX": max(x0, x1),
+                            "maxY": max(y0, y1)
+                        },
+                        "class_id": prediction["categories"][i],
+                        "domain": "pixel"
+                    }
+                    avg_scores[
+                        self.id_to_class[prediction["categories"][i]]
+                    ] += prediction["scores"][i]
+                    num_occurrences[self.id_to_class[prediction["categories"][i]]] += 1
+                    boxes.append(box)
+            else:
+                boxes = []
+
+            average_class_score = []
+            for cls in self.cats:
+                if cls["name"] not in num_occurrences:
+                    score = 0
+                else:
+                    score = avg_scores[cls["name"]] / num_occurrences[cls["name"]]
+                average_class_score.append(score)
+            result_table.add_data(
+                idx,
+                self.wandb.Image(val[1], boxes={
+                        "prediction": {
+                            "box_data": boxes,
+                            "class_labels": self.id_to_class
+                        }
+                    }
+                ),
+                *average_class_score
+            )
+
+        self.wandb.log({"val_results/result_table": result_table})
+
+    def save_checkpoint(self, save_dir, model_name, is_best, metadata=None):
         """
         Args:
             save_dir (str): save directory.
             model_name (str): model name.
             is_best (bool): whether the model is the best model.
+            metadata (dict): metadata to save corresponding to the checkpoint.
         """
+
+        if not self.log_checkpoints:
+            return
+
+        if "epoch" in metadata:
+            epoch = metadata["epoch"]
+        else:
+            epoch = None
+
         filename = os.path.join(save_dir, model_name + "_ckpt.pth")
         artifact = self.wandb.Artifact(
-            name=f"model-{self.run.id}",
-            type="model"
+            name=f"run_{self.run.id}_model",
+            type="model",
+            metadata=metadata
         )
         artifact.add_file(filename, name="model_ckpt.pth")
 
@@ -211,7 +354,23 @@ class WandbLogger(object):
         if is_best:
             aliases.append("best")
 
+        if epoch:
+            aliases.append(f"epoch-{epoch}")
+
         self.run.log_artifact(artifact, aliases=aliases)
 
     def finish(self):
         self.run.finish()
+
+    @classmethod
+    def initialize_wandb_logger(cls, args, exp, val_dataset):
+        wandb_params = dict()
+        prefix = "wandb-"
+        for k, v in zip(args.opts[0::2], args.opts[1::2]):
+            if k.startswith("wandb-"):
+                try:
+                    wandb_params.update({k[len(prefix):]: int(v)})
+                except ValueError:
+                    wandb_params.update({k[len(prefix):]: v})
+
+        return cls(config=vars(exp), val_dataset=val_dataset, **wandb_params)