glenn-jocher commited on
Commit
a600bae
·
unverified ·
1 Parent(s): a2d617e

Update model summary to display model name (#7101)

Browse files
Files changed (1) hide show
  1. 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
- LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
 
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)