瀏覽代碼

Add `on_fit_epoch_end` callback (#5232)

* Add `on_fit_epoch_end` callback

* Add results to train

* Update __init__.py
modifyDataloader
Glenn Jocher GitHub 2 年之前
父節點
當前提交
752d97df2c
沒有發現已知的金鑰在資料庫的簽署中 GPG Key ID: 4AEE18F83AFDEB23
共有 2 個文件被更改,包括 4 次插入2 次删除
  1. +3
    -1
      train.py
  2. +1
    -1
      utils/loggers/__init__.py

+ 3
- 1
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()

+ 1
- 1
utils/loggers/__init__.py 查看文件

@@ -131,7 +131,7 @@ class Loggers():
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

Loading…
取消
儲存