如何打印 model yolov5 model 的總結

Question

a.pt文件如何打印yolov5 model的model摘要？

# 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))

我試過 torchsummary、torchinfo 和 torchstat。 他們都沒有工作並且出錯了。 理想情況下，我想檢查 .network 中每一層的輸出/輸入維度。

Answer 1

您使用的代碼應該已經足夠了。

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)

這樣會打印出一個表格，顯示model中每一層的output維度，以及model的參數個數和memory的用法。

如果上面的代碼不充分或出錯，您可以執行以下操作以打印 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

從v6.0開始，官方model在運行models/yolo.py時打印輸入的維度。