Skip to main content

Module prune

Module prune 

Source
Expand description

Neural network pruning integration for Entrenar

This module provides training-time pruning capabilities that integrate with the Aprender pruning primitives. It implements:

  • Pruning Schedules: OneShot, Gradual, and Cubic sparsity schedules
  • Pruning Callback: Integration with the training loop
  • Calibration Pipeline: Activation collection for Wanda/SparseGPT
  • Prune-Finetune Pipeline: End-to-end pruning workflow

§Toyota Way Principles

  • Kaizen (Continuous Improvement): Gradual pruning schedules
  • Jidoka (Quality at Source): Validates sparsity at each step
  • Genchi Genbutsu (Go and See): Uses real activation data

§Example

use entrenar::prune::{PruningSchedule, PruningCallback, PruningConfig};
use aprender::pruning::{MagnitudeImportance, WandaPruner};

let schedule = PruningSchedule::Gradual {
    start_step: 1000,
    end_step: 10000,
    initial_sparsity: 0.0,
    final_sparsity: 0.5,
    frequency: 100,
};

let config = PruningConfig::default()
    .with_schedule(schedule)
    .with_target_sparsity(0.5);

let callback = PruningCallback::new(config);
trainer.add_callback(callback);

§References

  • Han, S., et al. (2015). Learning both weights and connections. NeurIPS.
  • Sun, M., et al. (2023). A simple and effective pruning approach. arXiv:2306.11695.
  • Zhu, M., & Gupta, S. (2017). To prune, or not to prune. arXiv:1710.01878.

Structs§

CalibrationCollector
Calibration data collector for pruning.
CalibrationConfig
Configuration for calibration data collection.
CalibrationDataConfig
Configuration for calibration data loading.
CalibrationDataLoader
Calibration data loader for pruning.
PruneFinetunePipeline
Prune-Finetune pipeline orchestrator.
PruneTrainer
Prune-finetune trainer that orchestrates the full pipeline.
PruneTrainerConfig
Configuration for the prune-finetune trainer.
PruningCallback
Callback for applying pruning during training.
PruningConfig
Configuration for pruning operations.
PruningMetrics
Metrics collected during pruning.

Enums§

PruneMethod
Pruning method selection.
PruningSchedule
Pruning schedule defining when sparsity increases during training.
PruningStage
Current stage of the pruning pipeline.
SparsityPatternConfig
Sparsity pattern selection.