1

How to print model summary of yolov5 model for a .pt file?

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu')

from torchstat import stat #try 1
stat(model, (3,640,640))

from torchsummary import summary #try 2
from torchinfo import summary #try 3
summary(model, (1,3,640,640))

I have tried torchsummary, torchinfo and torchstat. None of them work and errors out. Ideally, I want to check the output/input dimensions of every layer in the network.

desertnaut
  • 57,590
  • 26
  • 140
  • 166
user3303020
  • 933
  • 2
  • 12
  • 26
  • What are the errors you are getting? Which layer outputs do you want to print the shapes of? – Ivan Feb 22 '23 at 06:50

2 Answers2

0

The code you have used should have been sufficient.

from torchsummary import summary

# Create a YOLOv5 model
model = YOLOv5()

# Generate a summary of the model
input_size = (3, 640, 640)
summary(model, input_size=input_size)

This will print out a table that shows the output dimensions of each layer in the model, as well as the number of parameters and the memory usage of the model.

If above code wasn't sufficient or giving an error, you can do the following to print the dimensions of each layer in a YOLOv5 model.

import torch
from models.yolov5 import YOLOv5

# Create a YOLOv5 model
model = YOLOv5()

# Print the dimensions of each layer's inputs and outputs
for i, layer in enumerate(model.layers):
    print(f"Layer {i}: {layer.__class__.__name__}")
    x = torch.randn(1, 3, 640, 640)  # Create a random input tensor
    y = layer(x)
    print(f"\tInput dimensions: {x.shape}")
    print(f"\tOutput dimensions: {y.shape}")
Rajesh Kontham
  • 321
  • 1
  • 5
  • ``` summary[m_key]["input_shape"] = list(input[0].size()) AttributeError: 'list' object has no attribute 'size' ``` torchsummary doesn't work. – user3303020 Feb 26 '23 at 04:01
  • Isn't that the issue of "input[0].size()" rather than torchsummary? Why don't you try to print the shape (dimensions) of the image before you use it in torchsummary? It does mention " 'list' object has no attribute 'size' ". – Rajesh Kontham Feb 27 '23 at 01:11
0

From v6.0, the official model prints the dimensions of the inputs when running models/yolo.py.

desertnaut
  • 57,590
  • 26
  • 140
  • 166
Joey Park
  • 1