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:
parent
0bdd4ad379
commit
6391c60089
17 changed files with 76 additions and 47 deletions
|
|
@ -1,9 +1,10 @@
|
|||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from ultralytics.yolo.utils import LOGGER, RANK, TESTS_RUNNING, ops
|
||||
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
|
||||
from ultralytics.yolo.utils.torch_utils import model_info_for_loggers
|
||||
|
||||
try:
|
||||
import comet_ml
|
||||
|
|
@ -324,11 +325,7 @@ def on_fit_epoch_end(trainer):
|
|||
experiment.log_metrics(trainer.metrics, step=curr_step, epoch=curr_epoch)
|
||||
experiment.log_metrics(trainer.lr, step=curr_step, epoch=curr_epoch)
|
||||
if curr_epoch == 1:
|
||||
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), }
|
||||
experiment.log_metrics(model_info, step=curr_step, epoch=curr_epoch)
|
||||
experiment.log_metrics(model_info_for_loggers(trainer), step=curr_step, epoch=curr_epoch)
|
||||
|
||||
if not save_assets:
|
||||
return
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue