[Feat] Adding PruningCallback (#5618)

* wip * add pruning callback * add condition for duplicated weights * update on comments * update on comments * update on comments * add more tests * resolve flake8 * resolve on comments * update changelog * update on comments * update on comments * change order * remove ddp_spawn skip * update * typo * Update pytorch_lightning/callbacks/pruning.py Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> * Update pytorch_lightning/callbacks/pruning.py Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> * update on comments * forgot platform * update on comments * remove @rank_zero_only Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lightning-AI · Jan 27, 2021 · d0aaf98 · d0aaf98
1 parent 7e2e874
commit d0aaf98
Show file tree

Hide file tree

Showing 9 changed files with 631 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -59,6 +59,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `Recall` and `Precision` metrics (and their functional counterparts `recall` and `precision`) can now be generalized to Recall@K and Precision@K with the use of `top_k` parameter ([#4842](https://github.com/PyTorchLightning/pytorch-lightning/pull/4842))
 
 
+- Added `ModelPruning` Callback ([#5618](https://github.com/PyTorchLightning/pytorch-lightning/pull/5618))
+
+
 - Added `PyTorchProfiler` ([#5560](https://github.com/PyTorchLightning/pytorch-lightning/pull/5560))
 
 

diff --git a/docs/source/extensions/callbacks.rst b/docs/source/extensions/callbacks.rst
@@ -103,6 +103,7 @@ Lightning has a few built-in callbacks.
     LambdaCallback
     LearningRateMonitor
     ModelCheckpoint
+    ModelPruning
     ProgressBar
     ProgressBarBase
 

diff --git a/pytorch_lightning/callbacks/__init__.py b/pytorch_lightning/callbacks/__init__.py
@@ -20,6 +20,7 @@
 from pytorch_lightning.callbacks.lr_monitor import LearningRateMonitor
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
 from pytorch_lightning.callbacks.progress import ProgressBar, ProgressBarBase
+from pytorch_lightning.callbacks.pruning import ModelPruning
 
 __all__ = [
     'BackboneLambdaFinetuningCallback',
@@ -33,4 +34,5 @@
     'ModelCheckpoint',
     'ProgressBar',
     'ProgressBarBase',
+    'ModelPruning',
 ]
Original file line number	Diff line number	Diff line change
Expand Up		@@ -59,6 +59,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
		- `Recall` and `Precision` metrics (and their functional counterparts `recall` and `precision`) can now be generalized to Recall@K and Precision@K with the use of `top_k` parameter ([#4842](https://github.com/PyTorchLightning/pytorch-lightning/pull/4842))


		- Added `ModelPruning` Callback ([#5618](https://github.com/PyTorchLightning/pytorch-lightning/pull/5618))


		- Added `PyTorchProfiler` ([#5560](https://github.com/PyTorchLightning/pytorch-lightning/pull/5560))


Expand Down