mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-24 06:14:55 +08:00
fe27db2f6e
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com> Co-authored-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
30 lines
632 B
Python
30 lines
632 B
Python
# 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 {}
|
|
)
|