Schedules

When should I prune my network ?

The Schedule class allows you to create any schedule function according to 4 parameters:
- the sched_func: the function according to which the sparsity will evolve (e.g. linear, cos, …)
- the start_pct: the percentage of training steps at which the sparsification process will start
- the end_pct: the percentage of training steps at which the sparsification process will end
- the start_sparsity: the percentage of sparsity at which the model starts

One-Shot

The easiest schedule is the one-shot pruning, i.e. prune the network once. This can be done by simply returning the desired sparsity value. The moment when you want to prune will be controlled by the start_epoch argument in the SparsifyCallback.

one_shot.plot(50)

Iterative

Instead of pruning the network to desired sparsity in one step, you can do it iteratively. In fasterai, you can change the amount of iterations

iterative.plot(50)

To modify the default n_steps, you can use the partial function.

iterative = Schedule(partial(sched_iterative, n_steps=5), start_pct=0.2)
iterative.plot(50)

Automated Gradual Pruning

Some researchers have come up with more sophisticated schedules, such as the Automated Gradual Pruning.

agp.plot(50)

One-Cycle Pruning

one_cycle.plot(50)

On top of that, all of the schedules available in fastai by default are also available: - sched_cos - sched_linear

fig, ax = plt.subplots(1, 1, figsize=(8,5), dpi=100)
fig.patch.set_alpha(0.)
ax.patch.set_alpha(0.)

prune = np.linspace(0, 1, 1000)
sps = [cos([50], p) for p in prune]
plt.plot(prune, sps, c='#89d6c9', linewidth=2)
plt.xlabel('training iterations (Normalized)')
plt.ylabel('sparsity')

        
ax.spines['bottom'].set_color('#808080')
ax.spines['top'].set_color('#808080') 
ax.spines['right'].set_color('#808080')
ax.spines['left'].set_color('#808080')
ax.tick_params(axis='x', colors='#808080')
ax.tick_params(axis='y', colors='#808080')
ax.yaxis.label.set_color('#808080')
ax.xaxis.label.set_color('#808080')
fig, ax = plt.subplots(1, 1, figsize=(8,5), dpi=100)
fig.patch.set_alpha(0.)
ax.patch.set_alpha(0.)

prune = np.linspace(0, 1, 1000)
sps = [lin([50], p) for p in prune]
plt.plot(prune, sps, c='#89d6c9', linewidth=2)
plt.xlabel('training iterations (Normalized)')
plt.ylabel('sparsity')

        
ax.spines['bottom'].set_color('#808080')
ax.spines['top'].set_color('#808080') 
ax.spines['right'].set_color('#808080')
ax.spines['left'].set_color('#808080')
ax.tick_params(axis='x', colors='#808080')
ax.tick_params(axis='y', colors='#808080')
ax.yaxis.label.set_color('#808080')
ax.xaxis.label.set_color('#808080')

Dense-Sparse-Dense

You can also create even more interesting behaviours such as the DSD method, where you prune the model in the first place, then re-grow it to its initial amount of parameter.

dsd.plot(50)

See Also

  • SparsifyCallback - Apply sparsification during training using these schedules
  • PruneCallback - Apply structured pruning during training
  • Criteria - Different importance measures for selecting what to prune