From a600baed8efc6407ec4fb7a71fa1dbe3be23d441 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 22 Mar 2022 15:41:19 +0100 Subject: [PATCH] Update model summary to display model name (#7101) --- utils/torch_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index efcacc9..793c9c1 100644 --- 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)