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.
63 lines
1.9 KiB
63 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(trainer.args, name='parameters')
|
|
|
|
|
|
def on_batch_end(trainer):
|
|
train_loss = trainer.tloss
|
|
_log_scalers(trainer.label_loss_items(train_loss), "train", trainer.epoch)
|
|
|
|
|
|
def on_val_end(trainer):
|
|
metrics = trainer.metrics
|
|
val_losses = trainer.validator.loss
|
|
val_loss_dict = trainer.label_loss_items(val_losses)
|
|
_log_scalers(val_loss_dict, "val", trainer.epoch)
|
|
_log_scalers(metrics, "metrics", trainer.epoch)
|
|
|
|
if trainer.epoch == 0:
|
|
infer_speed = trainer.validator.speed[1]
|
|
model_info = {
|
|
"inference_speed": infer_speed,
|
|
"flops@640": get_flops(trainer.model),
|
|
"params": get_num_params(trainer.model)}
|
|
_log_scalers(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}
|