Add on_fit_epoch_end callback (ultralytics#5232)

* Add `on_fit_epoch_end` callback * Add results to train * Update __init__.py
fhkiel-mlaip · Oct 18, 2021 · dafbd67 · dafbd67
1 parent 724f355
commit dafbd67
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -423,8 +423,10 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
                                             plots=True,
                                             callbacks=callbacks,
                                             compute_loss=compute_loss)  # val best model with plots
+                    if is_coco:
+                        callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
 
-        callbacks.run('on_train_end', last, best, plots, epoch)
+        callbacks.run('on_train_end', last, best, plots, epoch, results)
         LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")
 
     torch.cuda.empty_cache()

diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
@@ -131,7 +131,7 @@ def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
             if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
                 self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
 
-    def on_train_end(self, last, best, plots, epoch):
+    def on_train_end(self, last, best, plots, epoch, results):
         # Callback runs on training end
         if plots:
             plot_results(file=self.save_dir / 'results.csv')  # save results.png