CODEWITHSUNDEEP
pytorchmodelyolov5
user3303020
user3303020

Reputation: 1043

how to print model summary of yolov5 model

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.

Upvotes: 1

Views: 2517

Answers (3)

Amirhosein Zarei
Amirhosein Zarei

Reputation: 1

try this way

from ultralytics import YOLO
from torchinfo import summary

model = YOLO("yolov5n.pt")  # Ensure you are using the correct model version
summary(model. Model, input_size=(1, 3, 640, 640))

https://github.com/ultralytics/yolov5/issues/11035#issuecomment-2249759900

Upvotes: 0

Joe Schmo
Joe Schmo

Reputation: 1

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

Upvotes: 0

Rajesh Kontham
Rajesh Kontham

Reputation: 349

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

Upvotes: 0

Related Questions