Tensorboard support (#435)

requirements.txt

@@ -7,3 +7,5 @@ torch >= 1.1.0
 matplotlib
 pycocotools
 tqdm
+tb-nightly
+future

train.py

@@ -59,7 +59,8 @@ def train(cfg,
           img_size=416,
           epochs=100,  # 500200 batches at bs 16, 117263 images = 273 epochs
           batch_size=16,
-          accumulate=4):  # effective bs = batch_size * accumulate = 16 * 4 = 64
+          accumulate=4,
+          write_to_tensorboard=False):  # effective bs = batch_size * accumulate = 16 * 4 = 64
     # Initialize
     init_seeds()
     weights = 'weights' + os.sep
@@ -226,7 +227,10 @@ def train(cfg,
 
             # Plot images with bounding boxes
             if epoch == 0 and i == 0:
-                plot_images(imgs=imgs, targets=targets, paths=paths, fname='train_batch%g.jpg' % i)
+                figure_image = plot_images(imgs=imgs, targets=targets, paths=paths, fname='train_batch%g.jpg' % i)
+
+                if write_to_tensorboard:
+                    tb_writer.add_image('train_batch', figure_image, dataformats='HWC')
 
             # Hyperparameter burn-in
             # n_burn = nb - 1  # min(nb // 5 + 1, 1000)  # number of burn-in batches
@@ -277,6 +281,21 @@ def train(cfg,
         with open('results.txt', 'a') as file:
             file.write(s + '%11.3g' * 7 % results + '\n')  # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
 
+        # Write Tensorboard results
+        if write_to_tensorboard:
+            tb_writer.add_scalar('GIoU/XY', mloss[0], epoch)
+            tb_writer.add_scalar('Width/Height', mloss[1], epoch)
+            tb_writer.add_scalar('Confidence', mloss[2], epoch)
+            tb_writer.add_scalar('Classification', mloss[3], epoch)
+            tb_writer.add_scalar('Train loss', mloss[4], epoch)
+            tb_writer.add_scalar('Precision', results[0], epoch)
+            tb_writer.add_scalar('Recall', results[1], epoch)
+            tb_writer.add_scalar('mAP', results[2], epoch)
+            tb_writer.add_scalar('F1', results[3], epoch)
+            tb_writer.add_scalar('Test loss GIoU', results[4], epoch)
+            tb_writer.add_scalar('Test loss obj', results[5], epoch)
+            tb_writer.add_scalar('Test loss cls', results[6], epoch)
+
         # Update best map
         fitness = results[2]  # mAP
         if fitness > best_fitness:
@@ -340,13 +359,29 @@ if __name__ == '__main__':
     opt = parser.parse_args()
     print(opt)
 
+
+
     if not opt.evolve:  # Train normally
+        # Tensorboard support, 
+        # start with "tensorboard --logdir=runs"
+        # go to localhost:6006
+        tensorboard_support = True
+        if version_to_tuple(torch.__version__)>= version_to_tuple("1.1.0"):
+            try:  
+                from torch.utils.tensorboard import SummaryWriter
+                tb_train_name = time.time()
+                print('here')
+                tb_writer = SummaryWriter('runs/{}'.format(tb_train_name))
+            except:
+                tensorboard_support = False
+
         results = train(opt.cfg,
                         opt.data,
                         img_size=opt.img_size,
                         epochs=opt.epochs,
                         batch_size=opt.batch_size,
-                        accumulate=opt.accumulate)
+                        accumulate=opt.accumulate,
+                        write_to_tensorboard=tensorboard_results)
 
     else:  # Evolve hyperparameters (optional)
         opt.notest = True  # only test final epoch

utils/utils.py

@@ -697,7 +697,9 @@ def plot_images(imgs, targets, paths=None, fname='images.jpg'):
             plt.title(s[:min(len(s), 40)], fontdict={'size': 8})  # limit to 40 characters
     fig.tight_layout()
     fig.savefig(fname, dpi=200)
+    fig_image = fig_to_data(fig)
     plt.close()
+    return fig_image
 
 
 def plot_test_txt():  # from utils.utils import *; plot_test()
@@ -809,3 +811,26 @@ def plot_results_orig(start=0, stop=0):  # from utils.utils import *; plot_resul
     fig.tight_layout()
     ax[4].legend()
     fig.savefig('results.png', dpi=200)
+
+
+def version_to_tuple(version):
+    # Used to compare versions of library
+    return tuple(map(int, (version.split("."))))
+
+
+def fig_to_data(fig):
+    # Used to convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
+    # param fig a matplotlib figure
+    # return a numpy 3D array of RGBA values
+
+    # draw the renderer
+    fig.canvas.draw()
+ 
+    # Get the RGBA buffer from the figure
+    w,h = fig.canvas.get_width_height()
+    buf = np.fromstring (fig.canvas.tostring_argb(), dtype=np.uint8)
+    buf.shape = (w, h, 4)
+ 
+    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
+    buf = np.roll(buf, 3, axis = 2)
+    return buf[:,:,:3] # Return RGB numpy image