New train profile argument for loggers (#2862)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-05-28 03:51:49 +02:00 committed by GitHub
parent 0bdd4ad379
commit 6391c60089
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 76 additions and 47 deletions

View file

@ -5,7 +5,7 @@ from time import time
from ultralytics.hub.utils import PREFIX, events
from ultralytics.yolo.utils import LOGGER
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
from ultralytics.yolo.utils.torch_utils import model_info_for_loggers
def on_pretrain_routine_end(trainer):
@ -24,11 +24,7 @@ def on_fit_epoch_end(trainer):
# Upload metrics after val end
all_plots = {**trainer.label_loss_items(trainer.tloss, prefix='train'), **trainer.metrics}
if trainer.epoch == 0:
model_info = {
'model/parameters': get_num_params(trainer.model),
'model/GFLOPs': round(get_flops(trainer.model), 3),
'model/speed(ms)': round(trainer.validator.speed['inference'], 3)}
all_plots = {**all_plots, **model_info}
all_plots = {**all_plots, **model_info_for_loggers(trainer)}
session.metrics_queue[trainer.epoch] = json.dumps(all_plots)
if time() - session.timers['metrics'] > session.rate_limits['metrics']:
session.upload_metrics()