From 95a1ba6eb24839a62e289ab92cdf59f4869f21ce Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 26 Jan 2021 21:16:01 -0800
Subject: [PATCH] W&B log epoch (#1946)

* W&B log epoch

* capitalize

* W&B log epoch

* capitalize

* Update train.py

New try using https://docs.wandb.ai/library/log#incremental-logging

* Update train.py

* Update test.py

* Update train.py

* Update plots.py

* Update train.py

* Update train.py

* label plot step -1

* update

* update

* update

* update

* update

* update

* Update train.py

* Update train.py
---
 test.py        | 4 ++--
 train.py       | 2 +-
 utils/plots.py | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/test.py b/test.py
index 891f6bef41c6..db344e722043 100644
--- a/test.py
+++ b/test.py
@@ -239,8 +239,8 @@ def test(data,
     if plots:
         confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
         if wandb and wandb.run:
-            wandb.log({"Images": wandb_images})
-            wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]})
+            val_batches = [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
+            wandb.log({"Images": wandb_images, "Validation": val_batches}, commit=False)
 
     # Save JSON
     if save_json and len(jdict):
diff --git a/train.py b/train.py
index 9e6bd8673726..5eff4bbac172 100644
--- a/train.py
+++ b/train.py
@@ -321,7 +321,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
                     #     tb_writer.add_graph(model, imgs)  # add model to tensorboard
                 elif plots and ni == 10 and wandb:
                     wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')
-                                           if x.exists()]})
+                                           if x.exists()]}, commit=False)
 
             # end batch ------------------------------------------------------------------------------------------------
         # end epoch ----------------------------------------------------------------------------------------------------
diff --git a/utils/plots.py b/utils/plots.py
index 4765069e0377..67f11bfd2011 100644
--- a/utils/plots.py
+++ b/utils/plots.py
@@ -295,7 +295,7 @@ def plot_labels(labels, save_dir=Path(''), loggers=None):
     # loggers
     for k, v in loggers.items() or {}:
         if k == 'wandb' and v:
-            v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]})
+            v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}, commit=False)
 
 
 def plot_evolution(yaml_file='data/hyp.finetune.yaml'):  # from utils.plots import *; plot_evolution()