From 9e859c949e20e84b96d6bfe6887f869c4acc03d5 Mon Sep 17 00:00:00 2001 From: Ivan Shcheklein Date: Sun, 11 Jun 2023 08:59:42 -0700 Subject: [PATCH] Use `model_info_for_loggers()` in DVC logger (#3127) --- ultralytics/yolo/utils/callbacks/dvc.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/ultralytics/yolo/utils/callbacks/dvc.py b/ultralytics/yolo/utils/callbacks/dvc.py index 4b30109..af5af16 100644 --- a/ultralytics/yolo/utils/callbacks/dvc.py +++ b/ultralytics/yolo/utils/callbacks/dvc.py @@ -4,7 +4,7 @@ import os import pkg_resources as pkg from ultralytics.yolo.utils import LOGGER, TESTS_RUNNING -from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params +from ultralytics.yolo.utils.torch_utils import model_info_for_loggers try: from importlib.metadata import version @@ -100,12 +100,7 @@ def on_fit_epoch_end(trainer): live.log_metric(metric, value) 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)} - - for metric, value in model_info.items(): + for metric, value in model_info_for_loggers(trainer).items(): live.log_metric(metric, value, plot=False) _log_plots(trainer.plots, 'train')