Spaces:
Running
Running
File size: 1,122 Bytes
e1832f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from ultralytics.utils import SETTINGS
try:
assert SETTINGS["raytune"] is True # verify integration is enabled
import ray
from ray import tune
from ray.air import session
except (ImportError, AssertionError):
tune = None
def on_fit_epoch_end(trainer):
"""
Sends training metrics to Ray Tune at end of each epoch.
This function checks if a Ray Tune session is active and reports the current training metrics along with the
epoch number to Ray Tune's session.
Args:
trainer (ultralytics.engine.trainer.BaseTrainer): The Ultralytics trainer object containing metrics and epochs.
Examples:
>>> # Called automatically by the Ultralytics training loop
>>> on_fit_epoch_end(trainer)
"""
if ray.train._internal.session.get_session(): # check if Ray Tune session is active
metrics = trainer.metrics
session.report({**metrics, **{"epoch": trainer.epoch + 1}})
callbacks = (
{
"on_fit_epoch_end": on_fit_epoch_end,
}
if tune
else {}
)
|