Aidge pruning API

Introduction

Aidge provides pruning functionalities that can be used during learning.

Basic example of training pipeline with pruning in Aidge (for a single epoch):

# Define model, setup learning...
model = ...
dataprovider = ...
opt = ...

# Setup pruning
# 1. Insert Prune nodes
aidge_pruning.insert_weight_prune_nodes(model)

# 2. Configure pruning
pruner = aidge_pruning.IterNonStructPruner()
pruning_rate = aidge_pruning.constant_pr(0.1)
pruner.set_pruning_rate_scheduler(pruning_rate)
pruner.set_data_and_masks(list(aidge_pruning.prune_data_and_masks(model)))

# 3. Initialize masks
pruner.update_masks()

scheduler = aidge_core.SequentialScheduler(model)
for i, (input, label) in enumerate(tqdm(dataprovider)):
    # Forward pass
    pred = scheduler.forward(data=[input])[0]
    # Reset the gradient
    opt.reset_grad()
    # Compute loss
    loss = aidge_learning.loss.CELoss(pred, label)
    # Compute accuracy (optional)
    acc = aidge_learning.metrics.Accuracy(pred, label, 1)[0]
    # Backward pass
    scheduler.backward()
    # Optimize the parameters
    opt.update()

    # Update masks
    pruner.update_masks()

Components

Pruners

The base pruner class is :class:` aidge_pruning.Pruner`:

class aidge_pruning.Pruner

__init__(self: aidge_pruning.aidge_pruning.Pruner) → None

data_and_masks(self: aidge_pruning.aidge_pruning.Pruner) → list[tuple[aidge_core.aidge_core.Tensor, aidge_core.aidge_core.Tensor]]

pruning_rate(self: aidge_pruning.aidge_pruning.Pruner) → float

pruning_rate_scheduler(self: aidge_pruning.aidge_pruning.Pruner) → Aidge::PRScheduler

set_data_and_masks(self: aidge_pruning.aidge_pruning.Pruner, arg0: collections.abc.Sequence[tuple[aidge_core.aidge_core.Tensor, aidge_core.aidge_core.Tensor]]) → None

set_pruning_rate_scheduler(self: aidge_pruning.aidge_pruning.Pruner, arg0: Aidge::PRScheduler) → None

update_masks(self: aidge_pruning.aidge_pruning.Pruner) → None

The available pruners are:

class aidge_pruning.IterNonStructPruner

__init__(self: aidge_pruning.aidge_pruning.IterNonStructPruner, delta: SupportsFloat | SupportsIndex = 0.0010000000474974513) → None

update_masks(self: aidge_pruning.aidge_pruning.IterNonStructPruner) → None

class aidge_pruning.RandomPruner

__init__(self: aidge_pruning.aidge_pruning.RandomPruner) → None

update_masks(self: aidge_pruning.aidge_pruning.RandomPruner) → None

Pruning rate scheduling

The base pruning rate scheduling class is aidge_pruning.PRScheduler:

class aidge_pruning.PRScheduler

__init__(*args, **kwargs)

pr_profiling(self: aidge_pruning.aidge_pruning.PRScheduler, arg0: SupportsInt | SupportsIndex) → std::vector<float, std::allocator<float> >

pruning_rate(self: aidge_pruning.aidge_pruning.PRScheduler) → float

step(self: aidge_pruning.aidge_pruning.PRScheduler) → int

update(self: aidge_pruning.aidge_pruning.PRScheduler) → None

The available pruning rate schedulers are:

aidge_pruning.constant_pr(initial_pr: SupportsFloat | SupportsIndex) → aidge_pruning.aidge_pruning.PRScheduler

aidge_pruning.step_pr(initial_pr: SupportsFloat | SupportsIndex, step_size: SupportsInt | SupportsIndex, gamma: SupportsFloat | SupportsIndex = 0.10000000149011612) → aidge_pruning.aidge_pruning.PRScheduler

Recipes

aidge_pruning.insert_weight_prune_nodes(graphview: aidge_core.aidge_core.GraphView) → None

aidge_pruning.prune_data_and_masks(graphview: aidge_core.aidge_core.GraphView) → list[tuple[aidge_core.aidge_core.Tensor, aidge_core.aidge_core.Tensor]]

Operators

aidge_pruning.Prune(name: str = '') → aidge_core.aidge_core.Node