Spaces:
Build error
Build error
File size: 2,598 Bytes
2283b14 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# Ultralytics YOLO 🚀, GPL-3.0 license
import json
from time import time
import torch
from ultralytics.hub.utils import PREFIX, sync_analytics
from ultralytics.yolo.utils import LOGGER
def on_pretrain_routine_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Start timer for upload rate limit
LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
session.t = {'metrics': time(), 'ckpt': time()} # start timer on self.rate_limit
def on_fit_epoch_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
session.metrics_queue[trainer.epoch] = json.dumps(trainer.metrics) # json string
if time() - session.t['metrics'] > session.rate_limits['metrics']:
session.upload_metrics()
session.t['metrics'] = time() # reset timer
session.metrics_queue = {} # reset queue
def on_model_save(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Upload checkpoints with rate limiting
is_best = trainer.best_fitness == trainer.fitness
if time() - session.t['ckpt'] > session.rate_limits['ckpt']:
LOGGER.info(f"{PREFIX}Uploading checkpoint {session.model_id}")
session.upload_model(trainer.epoch, trainer.last, is_best)
session.t['ckpt'] = time() # reset timer
def on_train_end(trainer):
session = getattr(trainer, 'hub_session', None)
if session:
# Upload final model and metrics with exponential standoff
LOGGER.info(f"{PREFIX}Training completed successfully ✅\n"
f"{PREFIX}Uploading final {session.model_id}")
session.upload_model(trainer.epoch, trainer.best, map=trainer.metrics['metrics/mAP50-95(B)'], final=True)
session.alive = False # stop heartbeats
LOGGER.info(f"{PREFIX}View model at https://hub.ultralytics.com/models/{session.model_id} 🚀")
def on_train_start(trainer):
sync_analytics(trainer.args)
def on_val_start(validator):
sync_analytics(validator.args)
def on_predict_start(predictor):
sync_analytics(predictor.args)
def on_export_start(exporter):
sync_analytics(exporter.args)
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,
"on_train_start": on_train_start,
"on_val_start": on_val_start,
"on_predict_start": on_predict_start,
"on_export_start": on_export_start}
|