You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
24 lines
608 B
24 lines
608 B
# Ultralytics YOLO 🚀, AGPL-3.0 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.""" |
|
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 {}
|
|
|