You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
56 lines
1.9 KiB
56 lines
1.9 KiB
from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params
|
|
|
|
try:
|
|
import clearml
|
|
from clearml import Task
|
|
|
|
assert hasattr(clearml, '__version__')
|
|
except (ImportError, AssertionError):
|
|
clearml = None
|
|
|
|
|
|
def _log_scalers(metric_dict, group="", step=0):
|
|
task = Task.current_task()
|
|
if task:
|
|
for k, v in metric_dict.items():
|
|
task.get_logger().report_scalar(group, k, v, step)
|
|
|
|
|
|
def before_train(trainer):
|
|
# TODO: reuse existing task
|
|
task = Task.init(project_name=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv5',
|
|
task_name=trainer.args.name if trainer.args.name != 'exp' else 'Training',
|
|
tags=['YOLOv5'],
|
|
output_uri=True,
|
|
reuse_last_task_id=False,
|
|
auto_connect_frameworks={'pytorch': False})
|
|
task.connect(dict(trainer.args), name='General')
|
|
|
|
|
|
def on_batch_end(trainer):
|
|
_log_scalers(trainer.label_loss_items(trainer.tloss, prefix="train"), "train", trainer.epoch)
|
|
|
|
|
|
def on_val_end(trainer):
|
|
_log_scalers(trainer.label_loss_items(trainer.validator.loss, prefix="val"), "val", trainer.epoch)
|
|
_log_scalers({k: v for k, v in trainer.metrics.items() if k.startswith("metrics")}, "metrics", trainer.epoch)
|
|
if trainer.epoch == 0:
|
|
model_info = {
|
|
"inference_speed": trainer.validator.speed[1],
|
|
"flops@640": get_flops(trainer.model),
|
|
"params": get_num_params(trainer.model)}
|
|
Task.current_task().connect(model_info, 'Model')
|
|
|
|
|
|
def on_train_end(trainer):
|
|
task = Task.current_task()
|
|
if task:
|
|
task.update_output_model(model_path=str(trainer.best), model_name='Best Model', auto_delete_file=False)
|
|
|
|
|
|
callbacks = {
|
|
"before_train": before_train,
|
|
"on_val_end": on_val_end,
|
|
"on_batch_end": on_batch_end,
|
|
"on_train_end": on_train_end}
|