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.572080 0.413692 0.847091 00:24 
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: 1252/1478 = 0.8471  Wilson 95% [0.8278, 0.8645]
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.331890 0.341828 0.868742 00:29
1 0.346365 0.327003 0.851150 00:28
2 0.287846 0.295884 0.893775 00:27
3 0.238572 0.515673 0.857916 00:23
4 0.179014 0.250834 0.893099 00:18
5 0.146992 0.480372 0.845737 00:14
6 0.102759 0.274859 0.913396 00:14
7 0.069884 0.218804 0.925575 00:12
8 0.040234 0.199090 0.928281 00:13
9 0.026111 0.203525 0.930311 00:13 
Pruning ratio at the end of epoch 0: 0.40%
Pruning ratio at the end of epoch 1: 1.57%
Pruning ratio at the end of epoch 2: 5.68%
Pruning ratio at the end of epoch 3: 16.06%
Pruning ratio at the end of epoch 4: 29.25%
Pruning ratio at the end of epoch 5: 36.69%
Pruning ratio at the end of epoch 6: 39.14%
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      224      258,048     
0.6.0.conv2                         Conv2d       224      203      409,248     
0.6.0.downsample.0                  Conv2d       128      203      25,984      
0.6.1.conv1                         Conv2d       203      18       32,886      
0.6.1.conv2                         Conv2d       18       203      32,886      
0.7.0.conv1                         Conv2d       203      1        1,827       
0.7.0.conv2                         Conv2d       1        512      4,608       
0.7.0.downsample.0                  Conv2d       203      512      103,936     
0.7.1.conv1                         Conv2d       512      1        4,608       
0.7.1.conv2                         Conv2d       1        512      4,608       
1.4                                 Linear       1024     502      514,048     
1.8                                 Linear       502      2        1,004       
-------------------------------------------------------------------------------------
Total                                                              2,074,843   
Original                                                           11,704,896  
Reduction                                                               82.27%

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.08 M parameters, 3349 conv output channels
conv output channels removed:  30.2%
parameter ratio:               0.18
MACs ratio at 224x224:         0.6335
MACs ratio at 64x64:           0.6346

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 → 3349, 30.2% 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 18 output channels in 0.6.1.conv1 and a single one in both 0.7.*.conv1 — and the head 1.4 goes from 512 to 502 units. Neither the parameter ratio (0.18) nor, in this run, the channel ratio (30.2%) 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: 1375/1478 = 0.9303  Wilson 95% [0.9162, 0.9422]

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.6335 at 224x224 against 0.6346 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