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

ultralytics 8.0.83 Neptune AI logging addition (#2130)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Snyk bot <snyk-bot@snyk.io>
Co-authored-by: Toutatis64 <Toutatis64@users.noreply.github.com>
Co-authored-by: M. Tolga Cangöz <46008593+standardAI@users.noreply.github.com>
Co-authored-by: Talia Bender <85292283+taliabender@users.noreply.github.com>
Co-authored-by: Ophélie Le Mentec <17216799+ouphi@users.noreply.github.com>
Co-authored-by: Kadir Şahin <68073829+ssahinnkadir@users.noreply.github.com>
Co-authored-by: Yonghye Kwon <developer.0hye@gmail.com>
This commit is contained in:
Glenn Jocher 2023-04-19 21:25:47 +02:00 • committed by GitHub
parent 55a03ad85f
commit 6c082ebd6f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 178 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

10
ultralytics/yolo/utils/callbacks/mlflow.py
View file

@ -52,19 +52,12 @@ def on_fit_epoch_end(trainer):
run.log_metrics(metrics=metrics_dict, step=trainer.epoch)
def on_model_save(trainer):
"""Logs model and metrics to mlflow on save."""
if mlflow:
run.log_artifact(trainer.last)
def on_train_end(trainer):
"""Called at end of train loop to log model artifact info."""
if mlflow:
root_dir = Path(__file__).resolve().parents[3]
run.log_artifact(trainer.last)
run.log_artifact(trainer.best)
model_uri = f'runs:/{run_id}/'
run.register_model(model_uri, experiment_name)
run.pyfunc.log_model(artifact_path=experiment_name,
code_path=[str(root_dir)],
artifacts={'model_path': str(trainer.save_dir)},
@ -74,5 +67,4 @@ def on_train_end(trainer):
callbacks = {
'on_pretrain_routine_end': on_pretrain_routine_end,
'on_fit_epoch_end': on_fit_epoch_end,
'on_model_save': on_model_save,
'on_train_end': on_train_end} if mlflow else {}