File size: 608 Bytes
ab854b9 |
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 |
# Ultralytics YOLO 🚀, AGPL-3.0 license
from ultralytics.utils import SETTINGS
try:
import ray
from ray import tune
from ray.air import session
assert SETTINGS['raytune'] is True # verify integration is enabled
except (ImportError, AssertionError):
tune = None
def on_fit_epoch_end(trainer):
"""Sends training metrics to Ray Tune at end of each epoch."""
if ray.tune.is_session_enabled():
metrics = trainer.metrics
metrics['epoch'] = trainer.epoch
session.report(metrics)
callbacks = {
'on_fit_epoch_end': on_fit_epoch_end, } if tune else {}
|