Update model summary to display model name (ultralytics#7101)

fhkiel-mlaip · Mar 22, 2022 · f5a84dc · f5a84dc
1 parent bb378e4
commit f5a84dc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -229,7 +229,8 @@ def model_info(model, verbose=False, img_size=640):
     except (ImportError, Exception):
         fs = ''
 
-    LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
+    name = model.yaml_file.rstrip('.yaml').replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model'
+    LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
 
 
 def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)