Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64445
AO Team is migrating the existing torch.quantization into torch.ao.quantization. We are doing it one file at a time to make sure that the internal callsites are updated properly.
This migrates the quantize.py from torch.quantization to torch.ao.quantization.
At this point both locations will be supported. Eventually the torch.quantization will be deprecated.
Test Plan: `buck test mode/dev //caffe2/test:quantization`
Reviewed By: HDCharles
Differential Revision: D30734870
fbshipit-source-id: dc204f3cc46bff2cc81c95159eab9d333b43bb4b
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64086
AO Team is migrating the existing torch.quantization into torch.ao.quantization. We are doing it one file at a time to make sure that the internal callsites are updated properly.
This migrates the `quantize.py` from torch.quantization to `torch.ao.quantization`.
At this point both locations will be supported. Eventually the torch.quantization will be deprecated.
Test Plan: `buck test mode/opt //caffe2/test:quantization`
Reviewed By: jerryzh168, raghuramank100
Differential Revision: D30055886
fbshipit-source-id: 8ef7470f9fa640c0042bef5bb843e7a05ecd0b9f
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63520
Rather than having to call `module.parametrizations.weight[0].pruned_outputs` each time we need to access the set of pruned indices, we add a getter `get_module_pruned_outputs` which takes the module as an argument and returns the set.
This is used for testing.
ghstack-source-id: 136561130
Test Plan:
` buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1N4gK
Reviewed By: z-a-f
Differential Revision: D30374558
fbshipit-source-id: e38dfee0879cadde52b942e899a3d8d7151ee493
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63519
If the pruner should be pruning biases along with weights, then if the model has BatchNorm2d following pruned Conv2d layers, then the corresponding channels of the BatchNorm must also be pruned.
Specifically, they need to zeroed out, rather than fully removed, since in eager mode, the dimensions between layers need to be preserved.
To do this, we add a pruning parametrization called `ZeroesParametrization` which zeroes out pruned channels, rather than removing them.
The user must provide in the config, a tuple of the Conv2d and BatchNorm layers that go together. The `prepare` method will add the tuple to the `module_groups`; then it will add a PruningParametrization to the Conv2d layer, and a ZeroesParametrization to BatchNorm, and then set their pruned sets to be the same set. That way, during `step`, both masks are updated with the same pruned indices.
ghstack-source-id: 136562278
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1N1P6
Reviewed By: z-a-f
Differential Revision: D30349855
fbshipit-source-id: 3199d3688d5a70963f9b32d7a8fdac3962ae6a65
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63202
By default, the prune will also prune biases, such that the whole output channel is removed. The user can manually set `also_prune_bias` to False when calling `prepare` if they don't want the bias to be pruned.
ghstack-source-id: 136466671
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1MV32
modify `fusion_tests` according to API change
`buck test mode/opt //scripts/kazhou:fusion_tests`
https://pxl.cl/1NbKz
Reviewed By: z-a-f
Differential Revision: D30294494
fbshipit-source-id: c84655648bee0035559195ca855b98fb7edaa134
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63178
Update base pruner API to match base sparsifier API as defined in D28970960 / PR58955
Changes include:
- `enable_mask_update = True` in `__init__`
- `prepare` takes model and config instead of constructor
- convert functionality renamed to `squash_mask`, `convert` method call now raises Error
- `activation_handles` ad `bias_handles` initialized in `_prepare` instead of constructor
ghstack-source-id: 136467595
Test Plan:
Function names updates according to changes
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1MTgH
TODO will need to modify `fbcode/scripts/kazhou/fusion_tests.py` to use new API
Reviewed By: z-a-f
Differential Revision: D30287179
fbshipit-source-id: d4727bea1873b500f2d4bb784db26d532bf26cce
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63158
Combined functionality for `ActivationReconstruction` for both Linear and Conv2d in one class. The only difference between the old classes was the size and indexing of the reconstructed tensor -- that logic can be generalized by iterating over the size of `output`.
ghstack-source-id: 136467465
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1MSSv
Reviewed By: raghuramank100
Differential Revision: D30282765
fbshipit-source-id: 08a1e4e0650511019fff85cf52b41dd818b0c7f8
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62430
The bias hook is a forward hook that is part of the pruning parametrization; it is attached after the activation reconstruction forward hook, so adding the bias occurs after zeros are reinserted to the pruned activation.
This diff/PR amends the bias hook to work for Conv2d layers, in addition to Linear layers. The reshaping of the ._bias parameter ensures that it is added to the right dimension of the output.
ghstack-source-id: 135097700
Test Plan:
Added tests for `Conv2dB()`, a model with Conv2d layers that have `bias=True`.
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1MfgL
Reviewed By: jerryzh168
Differential Revision: D29979571
fbshipit-source-id: c1a7e9fabc8b3c9d0050bd6b6c6a631ddfdf2a68
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61778
Adding Conv2d as supported modules for the pruner. Previously the pruner only supported Linear layers. This addition includes:
- adding a Conv2d activation reconstruction forward hook to match Conv2d weight shapes
- in `prepare`, checking the type of the module and using the corresponding activation forward hook
ghstack-source-id: 134143557
Test Plan:
Added conv2d tests
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1LLf3
Reviewed By: jerryzh168
Differential Revision: D29719045
fbshipit-source-id: 6a9f91b96992c552fff32f0e5a6e22f16eb7077b
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61592
Add activation handles for each layer (stored in a list), so they can each be removed.
We don't remove them in the `convert` in eager mode because we aren't modifying output/input layer dimensions. We will need this in Fx mode though.
ghstack-source-id: 133497376
Test Plan:
Added some tests to make sure `model(x)` runs without error.
`buck test mode/dev-nosan //caffe2/test:ao --
TestBasePruner`
https://pxl.cl/1LBf4
Reviewed By: z-a-f
Differential Revision: D29682789
fbshipit-source-id: 9185702736e5f7f4320754ffef441610738ac154
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61425
Adding handle for activation reconstruction and bias forward hooks so they can be removed later
ghstack-source-id: 133244536
Test Plan:
This change should not affect behavior yet, but to double check:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1LpM9
Reviewed By: z-a-f
Differential Revision: D29619720
fbshipit-source-id: c7428d2d0325cd11ce7919e0b67321e8cc196041
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59898
In `weight_norm_sparsifier`, the name of the argument `sparsity_pattern` is not intuitive for an argument describing the shape of the sparse block. It has been changed to `sparse_block_shape`.
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestWeightNormSparsifier`
https://pxl.cl/1LhRM
Reviewed By: z-a-f
Differential Revision: D29077045
fbshipit-source-id: 0cf9c5387d41ca8e839ee050d71f4fe477374143
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59771
Implements a specific sparsity scheduler, that uses a user-provided lambda's to change the levels.
Test Plan:
```
python test/test_ao_sparsity.py
```
Imported from OSS
Differential Revision:
D29070604
D29070604
Reviewed By: raghuramank100
Pulled By: z-a-f
fbshipit-source-id: c7ccbe63fe4cd6a0c3563541b7fcf93a99d0e62f
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59770
Implements the base scheduler class for changing the sparsity levels in the sparsifier.
Test Plan:
```
python test/test_ao_sparsity.py
```
Imported from OSS
Differential Revision:
D29070603
D29070603
Reviewed By: raghuramank100
Pulled By: z-a-f
fbshipit-source-id: 0b160e4eb0a2a303d2d19e6a3beb4784002b2cb7
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/58955
Implements the weight norm sparsifier.
This type of sparsifier computes the norm of the weights, sorts them, and zeroes-out the target fraction of them.
The main imeplemented method is `update_mask`, which holds the main logic of changing the masks.
Test Plan:
```
python test/test_ao_sparsity.py
```
Imported from OSS
Differential Revision:
D28970960
D28970960
Reviewed By: raghuramank100
Pulled By: z-a-f
fbshipit-source-id: 8f2a4360ad877f430cdc1065c6777106938b58d5
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/58707
Minor refactor that changes the format of the import.
This is done to avoid accidental circular dependencies.
Test Plan:
```
python test/test_ao_sparsity.py
```
Imported from OSS
Differential Revision:
D28970961
D28970961
Reviewed By: raghuramank100
Pulled By: z-a-f
fbshipit-source-id: c312742f5e218c435a1a643532f5842116bfcfff
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/58704
Implements the base sparsifier class based on the #59835 RFC documents.
This PR implements the base class for the sparsification. Specifically, the prepare method is implemented.
Test Plan:
```
python test/test_ao_sparsity.py
```
Imported from OSS
Differential Revision:
D28970958
D28970958
Reviewed By: raghuramank100
Pulled By: z-a-f
fbshipit-source-id: 0ef98a445c0a0aca22ce5708e34a9f94606d0e2b
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60395
Experimental folder so other developers know this is work in progress
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1KGJD
Reviewed By: z-a-f
Differential Revision: D29272319
fbshipit-source-id: 93eeeceba0376753efc9a5bb69a155278ceb2fca
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60278
Implemented `PruningParametrization`, which removes pruned rows, and `BasePruner`, which is the base class for structured pruning.
Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1KC2n
Reviewed By: z-a-f
Differential Revision: D29208349
fbshipit-source-id: f34e8e258bf13fa80292c2bd64d56f5ad1e72b6a
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/56617
This migrates the sparsity to the open source
Test Plan: `buck test mode/opt //caffe2/test:ao`
Reviewed By: raghuramank100
Differential Revision: D27812207
fbshipit-source-id: cc87d9d2b486269901a4ad9b483615741a1cd712