# Balance is Essence: Accelerating Sparse Training via Adaptive Gradient Correction

This is the code repository of RigL-based models in the following paper: Balance is Essence: Accelerating Sparse Training via Adaptive Gradient Correction.

Parts of this code repository is based on the following works by the machine learning community.

* https://github.com/varun19299/rigl-reproducibility


<span class="img_container center" style="display: block;">
    <p align="center">
    <img alt="RigL main image" height=350 src="img/dyn_sparse_train.png" style="display:block; margin-left: auto; margin-right: auto;" title="caption" />
    <br />
    <span class="img_caption" style="display: block; text-align: center;">Figure Courtsey: Evci et al. 2020.</span>
    </p>
</span>

## Getting Started

<details><summary><b>Install</b></summary>
<p>

* `python3.8`
* `pytorch`: 1.7.0+ (GPU support preferable).

Then,

* `make install`
</p>
</details>

<details><summary><b>W&B API key</b></summary>
<p>

Copy your WandB API key to `wandb_api.key`.
Will be used to login to your dashboard for visualisation. 
Alternatively, you can skip W&B visualisation, 
and set `wandb.use=False` while running the python code or `USE_WANDB=False` while running make commands.
</p>
</details>

<details><summary><b>Unit Tests</b></summary>
<p>

`make test`. Run `make help` to see specific make commands.
</p>
</details>

<details><summary><b>API Documentation</b></summary>
<p>

See `docs/index.html`, autogenerated via Sphinx.
Run `make docs.github` to refresh.
</p>
</details>

## Example Code

<details><summary><b>Train WideResNet-22-2 with RigL on CIFAR10</b></summary>
<p>

```
make cifar10.ERK.RigL DENSITY=0.2 SEED=0
````

Change `DENSITY` incase you want to use a different density (1 - sparsity) level.
See `outputs/CIFAR10/RigL_ERK/0.2/` for checkpoints etc. 
</p>
</details>

<details><summary><b>Evaluate WideResNet-22-2 with RigL on CIFAR10</b></summary>
<p>

Train WRN-22-2 with RigL as described above.
Place under `outputs/CIFAR10/RigL_ERK/0.2/+specific=cifar10_wrn_22_2_masking,seed=0`.

```
make cifar10.ERK.RigL DENSITY=0.2 SEED=0
```
</p>
</details>

## Main Results

<details><summary><b>Commands</b></summary>
<p>

The following make command runs all the main results described in our reproducibility report.

```
make cifar10 DENSITY=0.05,0.1,0.2,0.5
make cifar100 DENSITY=0.05,0.1,0.2,0.5
make cifar10_tune DENSITY=0.05,0.1,0.2,0.5
```

Use the `-n` flag to see which commands are executed.
Note that these runs are executed sequentially, although we include parallel processes for cifar10 runs of a particular method.
Eg: `cifar10.Random.RigL` runs RigL Random for densities `0.05,0.1,0.2,0.5`, `seed=0` in parallel.

It may be preferable to run specific make commands in parallel for this reason. See `make help` for an exhaustive list.

</p>
</details>

<details><summary><b>Visualization & Plotting Code</b></summary>
<p>

Run `make vis`.
</p>
</details>

