From b53d70ab9b75103a7bf2dda7b7d49201d34ad34b Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 22 Mar 2022 20:05:07 +0100 Subject: [PATCH] Model summary `pathlib` fix (#7104) Stems not working correctly for YOLOv5l with current .rstrip() implementation. After fix: ``` YOLOv5l summary: 468 layers, 46563709 parameters, 46563709 gradients, 109.3 GFLOPs ``` --- 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 793c9c184a44..72f8a0fd1659 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -11,6 +11,7 @@ import warnings from contextlib import contextmanager from copy import deepcopy +from pathlib import Path import torch import torch.distributed as dist @@ -229,7 +230,7 @@ def model_info(model, verbose=False, img_size=640): except (ImportError, Exception): fs = '' - name = model.yaml_file.rstrip('.yaml').replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model' + name = Path(model.yaml_file).stem.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}")