0
$\begingroup$

I trained a yolov8-medium object detection model, but I'm not sure how to apply pruning on it. What is the proper way to prune it?

I got this example code, but it only looks for conv1 layer which doesn't exist in the model, and throws an error:

import torch
from torch.nn.utils import prune
from ultralytics import YOLO

# Load your model
model = YOLO('best.pt')

# Display model layers
for name, module in model.named_modules():
    print(name)

# Specify the layer you want to prune and the amount
prune.l1_unstructured(model.model.conv1, name='weight', amount=0.5)

# To apply pruning permanently
prune.remove(model.model.conv1, 'weight')

The error:

    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'DetectionModel' object has no attribute 'conv1'

And here is the list of the layers it prints:

model
model.model
model.model.0
model.model.0.conv
model.model.0.bn
model.model.0.act
model.model.1
model.model.1.conv
model.model.1.bn
model.model.2
model.model.2.cv1
model.model.2.cv1.conv
model.model.2.cv1.bn
model.model.2.cv2
model.model.2.cv2.conv
model.model.2.cv2.bn
model.model.2.m
...
model.model.22.cv3.2.1.conv
model.model.22.cv3.2.1.bn
model.model.22.cv3.2.2
model.model.22.dfl
model.model.22.dfl.conv
$\endgroup$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.