Update model summary to display model name (#7101)
Browse files- utils/torch_utils.py +2 -1
utils/torch_utils.py
CHANGED
@@ -229,7 +229,8 @@ def model_info(model, verbose=False, img_size=640):
|
|
229 |
except (ImportError, Exception):
|
230 |
fs = ''
|
231 |
|
232 |
-
|
|
|
233 |
|
234 |
|
235 |
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|
|
|
229 |
except (ImportError, Exception):
|
230 |
fs = ''
|
231 |
|
232 |
+
name = model.yaml_file.rstrip('.yaml').replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model'
|
233 |
+
LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
234 |
|
235 |
|
236 |
def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
|