Pruned model showing same params

[tjasmin111]

Search before asking

• [X] I have searched the Ultralytics YOLO issues and discussions and found no similar questions.

Question

I was trying to prune a Yolov8 detection model, using this snippet. But when I load the model and try to detect, it still shows the same number of parameters (3,006,038 parameters). Is my pruning process correct? Anything missing here?

Snippet to prune:

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

# Load your model
model = YOLO('best.pt')
ratio = 0.5
numbers_to_check = [3, 4, 5, 6, 7, 8, 9, 10, 11]

for name, module in model.named_modules():
    print(name)
for name, module in model.named_modules():
    if ("cv1.conv" in name or "cv2.conv" in name) and any("."+str(num)+"." in name for num in numbers_to_check) and "m." not in name:
        print(f"Pruning layer (ratio: {ratio}): {name}")
        prune.l1_unstructured(module, name='weight', amount=ratio)
        prune.remove(module, 'weight')

# Save the pruned model
model.save(f'pruned_model_{ratio}.pt')

When I load the model, both models still show the same params numbers:

yolo detect predict model=pruned_best.pt source=dir
YOLOv8 summary (fused): 168 layers, 3,006,038 parameters, 0 gradients, 8.1 GFLOPs

Additional

No response

[Y-T-G]

Unstructured pruning doesn't change number of parameters.

https://stackoverflow.com/a/78301534/8061030

[tjasmin111]

I see thanks. So my snippet is correct and is the right way to do pruning? I need to fine tune it later too

[glenn-jocher]

@tjasmin111 yes, your snippet correctly applies unstructured pruning. After pruning, fine-tuning the model can help regain any lost accuracy.