Pruner

Structured pruning of a ResNet-18 - one shot, iterative, local, global, per layer

Pruner removes whole filters and channels from a model. The architecture itself gets smaller, so the pruned model runs on any hardware without sparse kernels.

import torch
from torchvision.models import resnet18, ResNet18_Weights
from fasterai.prune.pruner import Pruner
from fasterai.core.criteria import large_final

1. Remove 30% of the channels

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
before = sum(p.numel() for p in model.parameters())

print(model.conv1)
print(model.layer1[0].conv1)
print(f'{before:,} parameters')

Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
11,689,512 parameters

pruner = Pruner(model, pruning_ratio=0.3, context='local', criteria=large_final)
pruner.prune_model()

after = sum(p.numel() for p in model.parameters())
print(model.conv1)
print(model.layer1[0].conv1)
print(f'{after:,} parameters ({1 - after / before:.1%} fewer)')

Ignoring output layer: fc
Total ignored layers: 1
Conv2d(3, 44, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
Conv2d(44, 44, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
5,820,556 parameters (50.2% fewer)

pruning_ratio is a fraction of the channels: 0.3 takes conv1 from 64 to 44 output channels. The next convolution loses the matching input channels, which is why the parameter count falls by 50.2% and not by 30%.

criteria decides which filters are kept — see Criteria.

print_sparsity() lists the channel counts layer by layer.

pruner.print_sparsity()


Pruning Report:
-------------------------------------------------------------------------------------
Layer                               Type         In Ch    Out Ch   Params      
-------------------------------------------------------------------------------------
conv1                               Conv2d       3        44       6,468       
layer1.0.conv1                      Conv2d       44       44       17,424      
layer1.0.conv2                      Conv2d       44       44       17,424      
layer1.1.conv1                      Conv2d       44       44       17,424      
layer1.1.conv2                      Conv2d       44       44       17,424      
layer2.0.conv1                      Conv2d       44       89       35,244      
layer2.0.conv2                      Conv2d       89       89       71,289      
layer2.0.downsample.0               Conv2d       44       89       3,916       
layer2.1.conv1                      Conv2d       89       89       71,289      
layer2.1.conv2                      Conv2d       89       89       71,289      
layer3.0.conv1                      Conv2d       89       179      143,379     
layer3.0.conv2                      Conv2d       179      179      288,369     
layer3.0.downsample.0               Conv2d       89       179      15,931      
layer3.1.conv1                      Conv2d       179      179      288,369     
layer3.1.conv2                      Conv2d       179      179      288,369     
layer4.0.conv1                      Conv2d       179      358      576,738     
layer4.0.conv2                      Conv2d       358      358      1,153,476   
layer4.0.downsample.0               Conv2d       179      358      64,082      
layer4.1.conv1                      Conv2d       358      358      1,153,476   
layer4.1.conv2                      Conv2d       358      358      1,153,476   
fc                                  Linear       358      1000     359,000     
-------------------------------------------------------------------------------------
Total                                                              5,813,856   
Original                                                           11,689,512  
Reduction                                                               50.26%

print_sparsity totals the Conv2d and Linear parameters only (5,813,856) but compares that against every parameter of the original model (11,689,512), batch-norm included. Its 50.26% and the 50.2% printed above are therefore not the same ratio; the second one uses the same numerator and denominator.

2. Local and global context

context='local' applies the same ratio to every layer. context='global' compares filter importances across the whole network, so the ratio is met over the model as a whole and not layer by layer.

model_local = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
Pruner(model_local, 0.5, 'local', large_final).prune_model()

for i in range(1, 5):
    for j in (0, 1):
        conv = getattr(model_local, f'layer{i}')[j].conv1
        print(f'layer{i}.{j}.conv1: {conv.in_channels:3d} -> {conv.out_channels:3d} channels')
print(f'{sum(p.numel() for p in model_local.parameters()):,} parameters')

Ignoring output layer: fc
Total ignored layers: 1
layer1.0.conv1:  32 ->  32 channels
layer1.1.conv1:  32 ->  32 channels
layer2.0.conv1:  32 ->  64 channels
layer2.1.conv1:  64 ->  64 channels
layer3.0.conv1:  64 -> 128 channels
layer3.1.conv1: 128 -> 128 channels
layer4.0.conv1: 128 -> 256 channels
layer4.1.conv1: 256 -> 256 channels
3,055,880 parameters

model_global = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
Pruner(model_global, 0.5, 'global', large_final).prune_model()

for i in range(1, 5):
    for j in (0, 1):
        conv = getattr(model_global, f'layer{i}')[j].conv1
        print(f'layer{i}.{j}.conv1: {conv.in_channels:3d} -> {conv.out_channels:3d} channels')
print(f'{sum(p.numel() for p in model_global.parameters()):,} parameters')

Ignoring output layer: fc
Total ignored layers: 1
layer1.0.conv1:  64 ->  64 channels
layer1.1.conv1:  64 ->  64 channels
layer2.0.conv1:  64 -> 128 channels
layer2.1.conv1: 128 -> 126 channels
layer3.0.conv1: 128 -> 202 channels
layer3.1.conv1: 143 ->   9 channels
layer4.0.conv1: 143 -> 512 channels
layer4.1.conv1: 512 -> 512 channels
9,542,056 parameters

Both runs asked for 0.5, and removed it in different places. Local took half the output channels of every layer and ended at 3,055,880 parameters. Global compared all filters at once: it left the output channels of layer1, layer2 and layer4 where they were, took layer3.1.conv1 from 143 down to 9, and ended at 9,542,056 parameters. The input sides still moved with their producers — layer4.0.conv1 reads 143 channels instead of 256.

3. Iterative pruning

iterative_steps spreads the target ratio over several calls to prune_model().

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
pruner = Pruner(model, pruning_ratio=0.5, context='local', criteria=large_final, iterative_steps=5)

for step in range(5):
    pruner.prune_model()
    print(f'step {step + 1}: conv1 {model.conv1.out_channels:3d} channels, '
          f'{sum(p.numel() for p in model.parameters()):,} parameters')

Ignoring output layer: fc
Total ignored layers: 1
step 1: conv1  57 channels, 9,481,588 parameters
step 2: conv1  51 channels, 7,534,380 parameters
step 3: conv1  44 channels, 5,820,556 parameters
step 4: conv1  38 channels, 4,318,898 parameters
step 5: conv1  32 channels, 3,055,880 parameters

Five calls, five steps: conv1 goes 64 → 57 → 51 → 44 → 38 → 32 channels, and the last step lands on the same 3,055,880 parameters as the one-shot 50% run above.

4. One ratio per layer

pruning_ratio also accepts a {layer_name: fraction} dict. Dict targets need context='local', since a global context compares importances across layers.

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

per_layer = {
    'layer1.0.conv1': 0.2, 'layer1.0.conv2': 0.2,
    'layer2.0.conv1': 0.4, 'layer2.0.conv2': 0.4,
    'layer3.0.conv1': 0.6, 'layer3.0.conv2': 0.6,
    'layer4.0.conv1': 0.8, 'layer4.0.conv2': 0.8,
}
Pruner(model, per_layer, 'local', large_final).prune_model()

for i, ratio in zip(range(1, 5), (0.2, 0.4, 0.6, 0.8)):
    conv = getattr(model, f'layer{i}')[0].conv1
    print(f'layer{i}.0.conv1: {conv.out_channels:3d} channels ({ratio:.0%} removed)')

Ignoring output layer: fc
Total ignored layers: 1
Using per-layer pruning with 8 layer-specific ratios
layer1.0.conv1:  51 channels (20% removed)
layer2.0.conv1:  76 channels (40% removed)
layer3.0.conv1: 102 channels (60% removed)
layer4.0.conv1: 102 channels (80% removed)

Each named convolution loses its own share: 51 channels left of 64, 76 of 128, 102 of 256, 102 of 512. Layers absent from the dict keep all of theirs.

5. The pruned model still runs

Pruning rewires the dependencies between layers, so the model stays callable at the same input and output shapes.

model.eval()
x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    out = model(x)

print(f'input  {tuple(x.shape)}')
print(f'output {tuple(out.shape)}')

input  (1, 3, 224, 224)
output (1, 1000)

Summary

Call	What it gives you
`Pruner(model, 0.3, 'local', large_final)`	a pruner bound to this model; the ratio is a fraction of the channels, in [0, 1]
`pruner.prune_model()`	one pruning step, applied in place
`pruner.print_sparsity()`	per-layer input/output channel counts and the parameter total
`context='local'` / `'global'`	the same ratio in every layer, or one comparison across the whole network
`iterative_steps=n`	the target ratio spread over `n` calls to `prune_model()`
`pruning_ratio={name: fraction}`	one ratio per layer, with `context='local'`

Every figure on this page is a channel count or a parameter count on a torchvision ResNet-18 with its ImageNet weights. This page measures no accuracy and no latency.