HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add train loss and lr to loggers (#6732)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2023-12-01 22:09:20 +01:00 • committed by GitHub
parent 4425abce59
commit 9618025416
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 25 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

13
ultralytics/utils/callbacks/mlflow.py
View file

@ -32,7 +32,9 @@ try:
assert hasattr(mlflow, '__version__') # verify package is not directory
from pathlib import Path
PREFIX = colorstr('MLflow: ')
SANITIZE = lambda x: {k.replace('(', '').replace(')', ''): float(v) for k, v in x.items()}
except (ImportError, AssertionError):
mlflow = None
@ -81,11 +83,18 @@ def on_pretrain_routine_end(trainer):
f'{PREFIX}WARNING ⚠️ Not tracking this run')
def on_train_epoch_end(trainer):
"""Log training metrics at the end of each train epoch to MLflow."""
if mlflow:
mlflow.log_metrics(metrics=SANITIZE(trainer.label_loss_items(trainer.tloss, prefix='train')),
step=trainer.epoch)
mlflow.log_metrics(metrics=SANITIZE(trainer.lr), step=trainer.epoch)
def on_fit_epoch_end(trainer):
"""Log training metrics at the end of each fit epoch to MLflow."""
if mlflow:
sanitized_metrics = {k.replace('(', '').replace(')', ''): float(v) for k, v in trainer.metrics.items()}
mlflow.log_metrics(metrics=sanitized_metrics, step=trainer.epoch)
mlflow.log_metrics(metrics=SANITIZE(trainer.metrics), step=trainer.epoch)
def on_train_end(trainer):