Browse Source

Add `on_fit_epoch_end` callback (#5232)

* Add `on_fit_epoch_end` callback

* Add results to train

* Update __init__.py
modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
752d97df2c
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 2 deletions
  1. +3
    -1
      train.py
  2. +1
    -1
      utils/loggers/__init__.py

+ 3
- 1
train.py View File

plots=True, plots=True,
callbacks=callbacks, callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots 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)}") LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}")


torch.cuda.empty_cache() torch.cuda.empty_cache()

+ 1
- 1
utils/loggers/__init__.py View File

if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1: 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) 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 # Callback runs on training end
if plots: if plots:
plot_results(file=self.save_dir / 'results.csv') # save results.png plot_results(file=self.save_dir / 'results.csv') # save results.png

Loading…
Cancel
Save