Use `model_info_for_loggers()` in DVC logger (#3127)

single_channel
Ivan Shcheklein 1 year ago committed by GitHub
parent 8e60fc9276
commit 9e859c949e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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')

Loading…
Cancel
Save