xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add TensorBoard support (#87)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2022-12-24 14:37:46 +01:00
committed by GitHub
parent 248d54ca03
commit cb4f20f3cf
6 changed files with 133 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
ultralytics/yolo/utils/callbacks/clearml.py
View File

@ -9,47 +9,33 @@ 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):
def on_train_start(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'],
task = Task.init(project_name=trainer.args.project if trainer.args.project != 'runs/train' else 'YOLOv8',
task_name=trainer.args.name,
tags=['YOLOv8'],
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')
"Inference speed (ms/img)": round(trainer.validator.speed[1], 1),
"GFLOPs": round(get_flops(trainer.model), 1),
"Parameters": get_num_params(trainer.model)}
Task.current_task().connect(model_info, name='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)
Task.current_task().update_output_model(model_path=str(trainer.best),
model_name=trainer.args.name,
auto_delete_file=False)
callbacks = {
"before_train": before_train,
"on_train_start": on_train_start,
"on_val_end": on_val_end,
"on_batch_end": on_batch_end,
"on_train_end": on_train_end}
"on_train_end": on_train_end} if clearml else {}