how to print model summary of yolov5 model

Question

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.

Answer 1

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}")

Answer 2

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