0
$\begingroup$

I was trying to apply pruning to a YOLOv8 detection model. Following is my code and it runs successfully. I see the saved pruned model as well.

My YOLOv8 model has 22 layers (an example list of layers below) and I apply pruning to all cv1 and m.cv1 layers as shown in my code. My question is, which layer(s) should I choose to prune, and in what ratio? And should I prune both m and non-m modules such as these two?

model.model.15.cv1.conv
model.model.15.m.0.cv1.conv

And here is the code:

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

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

for name, module in model.named_modules():
    if "cv1.conv" in name:  # Check if the layer name contains "cv1"
        print(f"Pruning layer: {name}")
        prune.l1_unstructured(module, name='weight', amount=ratio)  # Prune the 'conv' submodule
        prune.remove(module, 'weight')  # Optional: Apply pruning permanently

# Save the pruned model
torch.save(model.state_dict(), 'pruned_model.pt')

enter image description here

$\endgroup$

0

Your Answer

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