YOLOv8

Iterative structured pruning of an ultralytics YOLOv8 detector

Pruner works on any nn.Module, including a detector coming from another library. This page prunes an ultralytics YOLOv8s on coco128, in ten steps, with a fine-tune after each step.

from copy import deepcopy
from datetime import datetime
from functools import partial
from pathlib import Path

import torch
import torch.nn as nn
import yaml
import torch_pruning as tp

from ultralytics import YOLO, __version__
from ultralytics.nn.modules import Detect, C2f, Conv, Bottleneck
from ultralytics.nn.tasks import attempt_load_one_weight
from ultralytics.engine.model import Model
from ultralytics.engine.trainer import BaseTrainer
from ultralytics.utils import LOGGER, RANK, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS
from ultralytics.utils.checks import check_yaml
from ultralytics.utils.torch_utils import initialize_weights, de_parallel

from fastai.vision.all import default_device
from fasterai.prune.all import Pruner
from fasterai.core.criteria import large_final
from fasterai.core.schedule import Schedule, sched_onecycle

1. Prepare the model

C2f concatenates the two halves of one convolution’s output, which ties channels that the pruner would have to split. replace_c2f_with_c2f_v2 rewrites those blocks into two separate convolutions carrying the same weights. A first fine-tune follows, so that the rest of the run starts from a model already fitted to coco128.

print('ultralytics', __version__, '| torch', torch.__version__)

model = YOLO('yolov8s.pt')
model.__setattr__("train_v2", train_v2.__get__(model))   # pruning-aware trainer, from the hidden cell

cfg = yaml_load(check_yaml('default.yaml'))
cfg.update(data='coco128.yaml', epochs=10, batch=8, verbose=False, amp=False, exist_ok=True)
batch_size = cfg['batch']

model.model.train()
replace_c2f_with_c2f_v2(model.model)
initialize_weights(model.model)
model.model = model.model.to(model.device).float()

validation_model = deepcopy(model)   # spare YOLO wrapper, used below to score pruned models

for p in model.model.parameters(): p.requires_grad = True
model.train_v2(pruning=True, **cfg)
print('fine-tuned, weights at', model.trainer.best)

ultralytics 8.3.162 | torch 2.9.1+cu128
fine-tuned, weights at runs/detect/train/weights/best.pt

The trainer hands the model back through attempt_load_one_weight, which returns the checkpoint with its gradients turned off. torch-pruning builds its dependency graph by tracing through autograd, so on a frozen model it finds nothing to prune and prune_model() returns without touching anything.

example_inputs = torch.randn(1, 3, cfg['imgsz'], cfg['imgsz']).to(model.device)

def size_of(net):
    "GMACs and M parameters of `net`"
    macs, params = tp.utils.count_ops_and_params(net, example_inputs)
    return f'{macs / 1e9:6.2f} GMACs  {params / 1e6:6.2f} M params'

frozen = deepcopy(model.model)                                     # exactly as the trainer returned it
heads = [m for m in frozen.modules() if isinstance(m, Detect)]
print('requires_grad on the reloaded model:', {p.requires_grad for p in frozen.parameters()})
print('before           ', size_of(frozen))

Pruner(frozen, 0.15, 'local', large_final, ignored_layers=heads).prune_model()
print('pruned, frozen   ', size_of(frozen))

frozen.requires_grad_(True)
Pruner(frozen, 0.15, 'local', large_final, ignored_layers=heads).prune_model()
print('pruned, grads on ', size_of(frozen))

requires_grad on the reloaded model: {False}
before             14.36 GMACs   11.17 M params
pruned, frozen     14.36 GMACs   11.17 M params
pruned, grads on   11.08 GMACs    8.39 M params

The frozen model comes back from prune_model() with exactly the parameters it went in with: nothing is raised, nothing is logged. The same model with requires_grad_(True) goes from 11.17 M to 8.39 M parameters and from 14.36 to 11.08 GMACs on one call. The loop below therefore re-enables the gradients before it builds its pruner.

2. Ten pruning steps

ignored_layers keeps the detection head out of the pruning, so the output shapes stay valid. iterative_steps and schedule split the 0.15 channel target over the ten calls to prune_model(): the schedule returns a progress in [0, 1] and the pruner multiplies the target by it.

Each step prints the MACs and parameters of the whole detector and the mAP50-95 before and after the fine-tune of that step.

ignored_layers = [m for m in model.model.modules() if isinstance(m, Detect)]
for p in model.model.parameters(): p.requires_grad = True    # see the cell above

pruner = Pruner(model.model, 0.15, 'local', large_final,
                ignored_layers=ignored_layers,
                iterative_steps=10,
                schedule=Schedule(partial(sched_onecycle, α=10, β=4)))

def val_map(net, name):
    "mAP50-95 of a DetectionModel, always through the same wrapper and the same validator config"
    cfg['name'], cfg['batch'] = name, 1
    validation_model.model.model = deepcopy(net.model)
    return validation_model.val(**cfg).box.map

# parity: the fine-tuned model scored through the shared wrapper and through a freshly
# loaded YOLO(best), to check the two give the same number before the table uses one of them
init_map = val_map(model.model, 'parity_shared')
probe = YOLO(model.trainer.best)
probe.model = deepcopy(model.model)
cfg['name'], cfg['batch'] = 'parity_reloaded', 1
print(f'parity: same model, two wrappers -> {init_map:.4f} and {probe.val(**cfg).box.map:.4f}')

base_macs, base_params = tp.utils.count_ops_and_params(model.model, example_inputs)

header = f"{'step':>4} {'GMACs':>8} {'M params':>9} {'% of base':>10} {'mAP pruned':>11} {'mAP tuned':>10}"
print()
print(header)
print('-' * len(header))
print(f"{0:>4} {base_macs / 1e9:>8.2f} {base_params / 1e6:>9.2f} {100:>9.1f}% "
      f"{'':>11} {init_map:>10.4f}")

for step in range(10):
    pruner.prune_model()
    macs, params = tp.utils.count_ops_and_params(pruner.model.to(default_device()),
                                                 example_inputs.to(default_device()))
    map_pruned = val_map(pruner.model, f'step_{step}_pre_val')

    cfg['name'], cfg['batch'] = f'step_{step}_finetune', batch_size
    model.model = pruner.model
    for p in model.model.parameters(): p.requires_grad = True
    model.train_v2(pruning=True, **cfg)
    best = model.trainer.best

    map_tuned = val_map(model.model, f'step_{step}_post_val')

    print(f"{step + 1:>4} {macs / 1e9:>8.2f} {params / 1e6:>9.2f} "
          f"{100 * params / base_params:>9.1f}% {map_pruned:>11.4f} {map_tuned:>10.4f}")

    torch.cuda.empty_cache()   # the deepcopies above keep the GPU busy between steps

    if init_map - map_tuned > 0.2:
        print('mAP dropped by more than 0.2 - stopping early')
        break

parity: same model, two wrappers -> 0.6484 and 0.6484

step    GMACs  M params  % of base  mAP pruned  mAP tuned
---------------------------------------------------------
   0    14.36     11.17     100.0%                 0.6484
   1    14.10     11.01      98.6%      0.6149     0.6728
   2    13.81     10.76      96.4%      0.6214     0.6708
   3    13.32     10.33      92.5%      0.6190     0.6805
   4    12.62      9.70      86.9%      0.5092     0.6694
   5    11.82      9.05      81.1%      0.4635     0.6413
   6    11.45      8.72      78.1%      0.4069     0.6396
   7    11.19      8.51      76.2%      0.5652     0.6471
   8    11.16      8.46      75.7%      0.6423     0.6670
   9    11.08      8.39      75.1%      0.6490     0.6718
  10    11.08      8.39      75.1%      0.6664     0.6741

Ten steps took the detector from 11.17 M to 8.39 M parameters and from 14.36 to 11.08 GMACs, 75.1% of what it started with. mAP pruned is the model right after a pruning step, mAP tuned after that step’s ten epochs; both columns go through the same wrapper and the same validator configuration, and the parity line above shows that wrapper agreeing with a freshly loaded YOLO(best) to the fourth decimal.

coco128 validates on its own training images, so both columns are a fit on the 128 images the step has just trained on. The pruned column dips to 0.4069 at step 6 and comes back; the tuned column stays between 0.6396 and 0.6805, and the last two steps move almost nothing — the one-cycle schedule has reached its target by then.

3. Check the model the loop left behind

best is the checkpoint the last fine-tune wrote. The validation below reuses cfg, so the validator runs with the same configuration as the table.

final = YOLO(best)
final_macs, final_params = tp.utils.count_ops_and_params(final.model, example_inputs.to(final.device))

print(f'{best}')
print(f'MACs   {base_macs / 1e9:.2f} G -> {final_macs / 1e9:.2f} G')
print(f'params {base_params / 1e6:.2f} M -> {final_params / 1e6:.2f} M')

runs/detect/step_9_finetune/weights/best.pt
MACs   14.36 G -> 11.08 G
params 11.17 M -> 8.39 M

cfg['name'], cfg['batch'] = 'final_val', 1
metrics = final.val(**cfg)          # the same validator configuration as the table above
print(f'mAP50-95 {metrics.box.map:.4f}   mAP50 {metrics.box.map50:.4f}')

mAP50-95 0.6741   mAP50 0.8714

4. Export

Under torch 2.9, ultralytics asks the ONNX exporter for an opset it can no longer produce, and the down-conversion fails; passing opset=18 avoids it.

onnx_path = final.export(format='onnx', opset=18)
print(onnx_path)

Applied 1 of general pattern rewrite rules.
runs/detect/step_9_finetune/weights/best.onnx

Measured on: ultralytics 8.3.162 and torch 2.9.1+cu128 (both printed above), YOLOv8s, coco128, batch 8 for the fine-tunes and 1 for validation, AMP off, a shared GPU, single run, one seed.

coco128.yaml uses images/train2017 as its validation split, so every mAP on this page is measured on the 128 images the model is fine-tuned on: a fit, not a generalisation number, and each row also carries the ten extra epochs its step added. Reloading the final checkpoint and validating it with the same cfg gives 0.6741, the last row of the table. MACs are a compute count at the model’s default image size; no inference time was measured.

Summary

Call	What it gives you
`Pruner(model.model, 0.15, 'local', large_final, ...)`	a pruner over the ultralytics `DetectionModel`; the ratio is a fraction of the channels
`ignored_layers=[m for m in ... if isinstance(m, Detect)]`	the detection head left untouched, so the output shapes stay valid
`iterative_steps=10`, `schedule=Schedule(...)`	the target spread over ten calls to `prune_model()`, following the schedule’s progress
`pruner.prune_model()`	one pruning step, applied in place
`tp.utils.count_ops_and_params(model, x)`	MACs and parameters of the current graph (torch-pruning)
`model.export(format='onnx', opset=18)`	ONNX export of the pruned detector

Two things this page needs that a torchvision model does not: C2f has to be rewritten into two separate convolutions before its channels can be pruned, and the parameters have to be set back to requires_grad=True after the trainer reloads the model from its checkpoint — the cell in section 1 shows what happens otherwise.

1. Prepare the model

2. Ten pruning steps

3. Check the model the loop left behind

4. Export

Summary

See Also