mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-29 02:14:22 +08:00
Use model_info_for_loggers()
in DVC logger (#3127)
This commit is contained in:
parent
8e60fc9276
commit
9e859c949e
@ -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…
x
Reference in New Issue
Block a user