Prune Callback

Prune filters during fastai training with PruneCallback 
from fastai.vision.all import *
from fasterai.prune.all import *
import torch_pruning as tp

Prune while training

PruneCallback removes filters during a fastai training run: on every training step it calls its Pruner, which removes the filters criteria ranks lowest, up to the ratio schedule prescribes at that point. The model that comes out of fit is structurally smaller, and training continues on it as the ratio grows.

1. Data and baseline

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64))

A pretrained ResNet-18, warmed up for one epoch on the cat/dog task:

learn = vision_learner(dls, resnet18, metrics=accuracy)
learn.unfreeze()
learn.fit_one_cycle(1)
epoch train_loss valid_loss accuracy time
0 0.601222 0.380172 0.845061 00:05 
def report(name, k, n, z=1.96):
    "A count with its n and its Wilson 95% interval"
    p, d, centre = k/n, 1 + z**2/n, k/n + z**2/(2*n)
    half = z*((p*(1-p) + z**2/(4*n))/n)**0.5
    print(f"{name}: {k}/{n} = {p:.4f}  Wilson 95% [{(centre-half)/d:.4f}, {(centre+half)/d:.4f}]")

def correct(learn):
    "Correct predictions and n on the validation set"
    n = len(learn.dls.valid_ds)
    with learn.no_bar(): acc = float(learn.validate()[1])
    return round(acc*n), n

report("after the warm-up epoch", *correct(learn))
after the warm-up epoch: 1249/1478 = 0.8451  Wilson 95% [0.8257, 0.8626]
sample_224 = torch.randn(1,3,224,224).to(default_device())
sample_64  = torch.randn(1,3,64,64).to(default_device())

base_macs, base_params = tp.utils.count_ops_and_params(learn.model, sample_224)
base_macs_64, _ = tp.utils.count_ops_and_params(learn.model, sample_64)
base_channels = sum(m.out_channels for m in learn.model.modules() if isinstance(m, nn.Conv2d))
print(f"{base_macs/1e9:.2f} GMACs at 224x224, {base_params/1e6:.2f} M parameters, {base_channels} conv output channels")
1.82 GMACs at 224x224, 11.70 M parameters, 4800 conv output channels

2. Training with PruneCallback

The call below asks the pruner for a global ratio of 0.4 over ten epochs:

  • pruning_ratio=0.4 - the ratio asked of the pruner, a fraction in [0, 1] (0.4 = 40%)
  • context='global' - compare filter importance across the whole model
  • criteria=large_final - keep the filters with the largest final weights
  • schedule=one_cycle - how the ratio grows from 0 to its target during training
pr_cb = PruneCallback(pruning_ratio=0.4, context='global', criteria=large_final, schedule=one_cycle)
learn.fit_one_cycle(10, cbs=pr_cb)
Ignoring output layer: 1.8
Total ignored layers: 1
epoch train_loss valid_loss accuracy time
0 0.354834 0.316359 0.871448 00:06
1 0.289825 0.298112 0.886333 00:06
2 0.245730 0.237911 0.900541 00:07
3 0.228931 0.332760 0.872801 00:07
4 0.184838 0.218940 0.914073 00:08
5 0.213857 0.227616 0.915426 00:07
6 0.266128 0.235636 0.897835 00:08
7 0.241007 0.234197 0.904601 00:06
8 0.230154 0.237446 0.900541 00:05
9 0.226757 0.228660 0.909337 00:05 
Pruning ratio at the end of epoch 0: 0.39%
Pruning ratio at the end of epoch 1: 1.54%
Pruning ratio at the end of epoch 2: 5.60%
Pruning ratio at the end of epoch 3: 15.91%
Pruning ratio at the end of epoch 4: 29.13%
Pruning ratio at the end of epoch 5: 36.64%
Pruning ratio at the end of epoch 6: 39.12%
Pruning ratio at the end of epoch 7: 39.79%
Pruning ratio at the end of epoch 8: 39.96%
Pruning ratio at the end of epoch 9: 40.00%

The callback keeps its Pruner as pr_cb.pruner, which reports the channel count of every convolution left in the model:

pr_cb.pruner.print_sparsity()

Pruning Report:
-------------------------------------------------------------------------------------
Layer                               Type         In Ch    Out Ch   Params      
-------------------------------------------------------------------------------------
0.0                                 Conv2d       3        64       9,408       
0.4.0.conv1                         Conv2d       64       64       36,864      
0.4.0.conv2                         Conv2d       64       64       36,864      
0.4.1.conv1                         Conv2d       64       64       36,864      
0.4.1.conv2                         Conv2d       64       64       36,864      
0.5.0.conv1                         Conv2d       64       128      73,728      
0.5.0.conv2                         Conv2d       128      128      147,456     
0.5.0.downsample.0                  Conv2d       64       128      8,192       
0.5.1.conv1                         Conv2d       128      128      147,456     
0.5.1.conv2                         Conv2d       128      128      147,456     
0.6.0.conv1                         Conv2d       128      239      275,328     
0.6.0.conv2                         Conv2d       239      183      393,633     
0.6.0.downsample.0                  Conv2d       128      183      23,424      
0.6.1.conv1                         Conv2d       183      19       31,293      
0.6.1.conv2                         Conv2d       19       183      31,293      
0.7.0.conv1                         Conv2d       183      1        1,647       
0.7.0.conv2                         Conv2d       1        512      4,608       
0.7.0.downsample.0                  Conv2d       183      512      93,696      
0.7.1.conv1                         Conv2d       512      1        4,608       
0.7.1.conv2                         Conv2d       1        512      4,608       
1.4                                 Linear       1024     506      518,144     
1.8                                 Linear       506      2        1,012       
-------------------------------------------------------------------------------------
Total                                                              2,064,446   
Original                                                           11,704,896  
Reduction                                                               82.36%

3. What the pruned model costs

pruned_macs, pruned_params = tp.utils.count_ops_and_params(learn.model, sample_224)
pruned_macs_64, _ = tp.utils.count_ops_and_params(learn.model, sample_64)
pruned_channels = sum(m.out_channels for m in learn.model.modules() if isinstance(m, nn.Conv2d))

print(f"{pruned_macs/1e9:.2f} GMACs at 224x224, {pruned_params/1e6:.2f} M parameters, {pruned_channels} conv output channels")
print(f"conv output channels removed:  {1 - pruned_channels/base_channels:.1%}")
print(f"parameter ratio:               {pruned_params/base_params:.2f}")
print(f"MACs ratio at 224x224:         {pruned_macs/base_macs:.4f}")
print(f"MACs ratio at 64x64:           {pruned_macs_64/base_macs_64:.4f}")
1.15 GMACs at 224x224, 2.07 M parameters, 3305 conv output channels
conv output channels removed:  31.1%
parameter ratio:               0.18
MACs ratio at 224x224:         0.6328
MACs ratio at 64x64:           0.6339

pruning_ratio=0.4 asks the pruner for a global ratio of 0.4; the cells above show what that produced here. Counting the output channels of every convolution, 4800 → 3305, 31.1% removed. In the report, 0.4 to 0.7 are module paths and not ratios: the four ResNet stages inside the fastai body, two blocks each. The stem 0.0, 0.4 and 0.5 keep their 64 and 128 channels, the removal falls on 0.6 and 0.7 — down to 19 output channels in 0.6.1.conv1 and a single one in both 0.7.*.conv1 — and the head 1.4 goes from 1024 to 506 units. Neither the parameter ratio (0.18) nor, in this run, the channel ratio (31.1%) equals the 0.4 asked for: the ratio is what torch-pruning targets over its prunable groups, and a group cannot go below one channel — two convolutions here sit at that bound.

report("after pruning", *correct(learn))
after pruning: 1344/1478 = 0.9093  Wilson 95% [0.8936, 0.9229]

Scope. ResNet-18 pretrained on ImageNet, PETS cat/dog labels, images resized to 64 px; one warm-up epoch, then ten epochs with PruneCallback; single run, no seed fixed. Both accuracies are measured on the 1478 validation images and printed by report with their Wilson 95% interval, but they come from different training budgets — one epoch against eleven — so they are not a dense-versus-pruned comparison. The absolute GMACs are for a 224x224 batch-1 input rather than the 64 px training resolution; the MACs ratio barely moves with that choice, 0.6328 at 224x224 against 0.6339 at 64x64, as the cell above prints. This page measures no latency.

Summary

Tool What it gives you
PruneCallback(pruning_ratio, context, criteria, schedule) Filters removed during training, following the schedule
pr_cb.pruner The Pruner the callback built, still usable after fit
Pruner.print_sparsity() Per-layer channel counts and the total parameter reduction
tp.utils.count_ops_and_params MACs and parameters of a model, at a given input size

pruning_ratio is a fraction in [0, 1] (0.4 = 40%). criteria and schedule take the objects exported by fasterai.core.criteria (large_final, small_final, random, movement, wanda, …) and fasterai.core.schedule (one_shot, iterative, agp, one_cycle, cos, lin, dsd).

See Also