Metadata-Version: 2.1
Name: autosparsity
Version: 1.0
Summary: Rockchip AutoSparsity Tools
Author: Rockchip
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Indenpendent
License-File: LICENSE
Requires-Dist: torch>=1.10.0
Requires-Dist: torchvision>=0.11.1

## 介绍

该工具用于torch模型训练阶段的权重自动稀疏化，已有的稀疏化方式包括：4:2输入方向稀疏，4:2输出方向稀疏, 16:4输入输出双向组合稀疏。
另外，稀疏方式支持用户自定义。
（权重shape为CoutCinKhKw, 输入方向即Cin, 输出方向即Cout）

| 稀疏化方式 | 稀疏率　| 说明 |
| :---- | :---- | :---- |
| 4:2输入方向稀疏 | 50% | 沿输入方向对权重稀疏化，refer to NVIDIA paper [《Accelerating sparse deep neural networks》](https://arxiv.org/abs/2104.08378)|
| 4:2输出方向稀疏 | 50% | 沿输出方向对权重稀疏化 |
| 16:4输入输出双向组合稀疏 | 75% | 先沿输入方向稀疏4:2，再沿输出方向稀疏4:2（或反之），前置训练模型可以是dense模型也可以是稀疏模型，建议使用单方向稀疏模型作为前置预训练模型，可减少稀疏化后的精度损失 |

## 生成wheel包

### step1

生成permutation_search_cuda.cpython-[python version]m-x86_64-linux-gnu.so。
```bash
cd autosparsity/permutation_tests/
bash run.sh
bash test.sh
```

### step2

打包wheel，在对应python平台编译打包。
```bash
cd -
python setup.py -py cp36
```

## 安装

安装对应python平台的包
```bash
pip install autosparsity-1.0-cp36-cp36m-linux_x86_64.whl
```

## 测试用例

该用例对torchvision中的resnet50做50%稀疏化。
```bash
python example/test.py
```

## 具体使用方法介绍

在训练代码中, 创建model之后，加入以下稀疏化代码：
```python
# insert model autosparsity code before training
import torch
import torchvision.models as models
from autosparsity.sparsity import sparsity_model

model = models.resnet34(pretrained=True).cuda()
mode = 0
sparsity_model(model, optimizer, mode)

# normal training
x, y = DataLoader(args)
for epoch in range(epochs):
    y_pred = model(x)
    loss = loss_func(y_pred, y)
    loss.backward()
    optimizer.step()
    ...
```

sparsity_model函数各参数说明如下

| 参数 | 说明 |
| :---- | ---- |
| model | 原训练模型 |
| optimizer | 原优化器，默认None |
| mode | 稀疏模式(0-3, 0:50%输入稀疏, 1:50%输出稀疏, 2:75%输入输出组合稀疏, 3:75%输出输入组合稀疏)，默认0 |
| verbose | log等级(0-2, 0:Errors, 1:Errors and Warnings, 2:Errors, warnings and info) ，默认2 |
| whitelist | 稀疏化支持的module列表，支持1d conv，2d conv，3d conv, linear, MultiheadAttention，默认[torch.nn.Linear, torch.nn.Conv2d] |
| allowed_layer_names | 允许稀疏化的层名，当用户配置时则只稀疏允许的层，默认None |
| disallowed_layer_names | 不允许稀疏化的层名，当用户配置时则会跳过该层，默认[] |
| fast | 设为True代表使用快速方法计算mask，默认为False(默认的mask计算方法针对部分模型会失效，若稀疏化报错可尝试将该参数设为True) |

**Attention* 模型稀疏化训练依赖于一个已训练好的前置模型。所以在做稀疏化训练之前，请先训练一个效果‘不错’的模型，在该模型基础上进行稀疏化训练。
