New train profile argument for loggers (#2862)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/utils/torch_utils.py

@@ -192,6 +192,29 @@ def get_num_gradients(model):
     return sum(x.numel() for x in model.parameters() if x.requires_grad)
 
 
+def model_info_for_loggers(trainer):
+    """
+    Return model info dict with useful model information.
+
+    Example for YOLOv8n:
+        {'model/parameters': 3151904,
+         'model/GFLOPs': 8.746,
+         'model/speed_ONNX(ms)': 41.244,
+         'model/speed_TensorRT(ms)': 3.211,
+         'model/speed_PyTorch(ms)': 18.755}
+    """
+    if trainer.args.profile:  # profile ONNX and TensorRT times
+        from ultralytics.yolo.utils.benchmarks import ProfileModels
+        results = ProfileModels([trainer.last], device=trainer.device).profile()[0]
+        results.pop('model/name')
+    else:  # only return PyTorch times from most recent validation
+        results = {
+            'model/parameters': get_num_params(trainer.model),
+            'model/GFLOPs': round(get_flops(trainer.model), 3)}
+    results['model/speed_PyTorch(ms)'] = round(trainer.validator.speed['inference'], 3)
+    return results
+
+
 def get_flops(model, imgsz=640):
     """Return a YOLO model's FLOPs."""
     try: