glenn-jocher
commited on
Model summary `pathlib` fix (#7104)
Browse filesStems 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 +2 -1
utils/torch_utils.py
CHANGED
@@ -11,6 +11,7 @@ import time
|
|
11 |
import warnings
|
12 |
from contextlib import contextmanager
|
13 |
from copy import deepcopy
|
|
|
14 |
|
15 |
import torch
|
16 |
import torch.distributed as dist
|
@@ -229,7 +230,7 @@ def model_info(model, verbose=False, img_size=640):
|
|
229 |
except (ImportError, Exception):
|
230 |
fs = ''
|
231 |
|
232 |
-
name = model.yaml_file.
|
233 |
LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
234 |
|
235 |
|
|
|
11 |
import warnings
|
12 |
from contextlib import contextmanager
|
13 |
from copy import deepcopy
|
14 |
+
from pathlib import Path
|
15 |
|
16 |
import torch
|
17 |
import torch.distributed as dist
|
|
|
230 |
except (ImportError, Exception):
|
231 |
fs = ''
|
232 |
|
233 |
+
name = Path(model.yaml_file).stem.replace('yolov5', 'YOLOv5') if hasattr(model, 'yaml_file') else 'Model'
|
234 |
LOGGER.info(f"{name} summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
|
235 |
|
236 |
|