Metadata-Version: 2.4
Name: rasnatune
Version: 0.0.3
Summary: Runtime compression tuning utilities for PyTorch
Author: rasnatune contributors
License-Expression: MIT
Keywords: pytorch,compression,sparsification,quantization
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Dynamic: license-file

# rasnatune

`rasnatune` は、ハードウェア寄りの量子化やスパース化を PyTorch モデルで学習・検証するための軽量ツールキットです。
既存モデルに一時的な forward hook を直接取り付けるため、モデルオブジェクトそのものや `state_dict()` のキーは変わりません。

## インストール

```bash
pip install rasnatune
```

## クイックスタート

```python
import torch
import torchvision

import rasnatune
from rasnatune import QuantizeIntermidiate, SparseWeight, skip_first_layer

model = torchvision.models.resnet50()

rasnatune.apply(
    model,
    SparseWeight(sparsity=0.5),
    targets=[torch.nn.Conv2d],
    filters=[skip_first_layer],
)
rasnatune.apply(
    model,
    QuantizeIntermidiate(ab=8, wb=8, ib=17),
    targets=[torch.nn.Conv2d, torch.nn.Linear],
    filters=[],
)

# 以後は通常の学習を行います。

rasnatune.apply(
    model,
    SparseWeight(sparsity=0.8),
    targets=[torch.nn.Conv2d],
    filters=[skip_first_layer],
)
```

## Apply

`apply()` は設定済み compressor インスタンスをモデルへ適用し、同じモデルインスタンスを返します。
渡された compressor は対象 module ごとに内部で複製されます。

```python
from rasnatune import apply, QuantizeIntermidiate, SparseWeight, skip_first_layer

apply(
    model,
    QuantizeIntermidiate(ab=8, wb=8, ib=17),
    targets=[torch.nn.Conv2d, torch.nn.Linear],
    filters=[],
)
apply(
    model,
    SparseWeight(sparsity=0.5),
    targets=[torch.nn.Conv2d],
    filters=[skip_first_layer],
)
```

`targets=` には module 型の list / tuple を渡します。
`filters=` には predicate を渡せます。
`skip_first_layer` は、`targets=` で選ばれた最初の module を除外します。

```python
apply(
    model,
    QuantizeIntermidiate(ab=8, wb=8, ib=17),
    targets=[torch.nn.Linear],
    filters=[],
)
```

同じ module に同じ種類の compressor を再適用すると、古い compressor は detach され、新しい compressor に差し替わります。
例えば学習途中で `SparseWeight(sparsity=0.5)` 相当から `sparsity=0.8` へ変更すると、該当 Conv2d のスパース率だけが更新されます。
`SparseWeight` と `QuantizeIntermidiate` は別種類なので、同じ module 上で共存できます。

## 公開 API

トップレベル API:

- `rasnatune.apply(model, compressor, targets=..., filters=...)`
- `rasnatune.remove_compression(model)`
- `rasnatune.compression_count(model)`
- `rasnatune.skip_first_layer`

compressor クラス:

- `rasnatune.Compressor`
- `rasnatune.QuantizeWeight`
- `rasnatune.QuantizeActivation`
- `rasnatune.QuantizeIntermidiate`
- `rasnatune.SparseWeight`
- `rasnatune.SparseActivation`

## AWI 量子化シミュレーション

`QuantizeIntermidiate` は、実際の PyTorch 実行は FP32 のまま保ちつつ、Activation / Weight / Intermediate の量子化挙動をシミュレートします。

- activation と weight は straight-through gradient 付きで量子化されます
- weight のスパース化は `SparseWeight` を同じ module に重ねて表現します
- `Linear` / `Conv2d` の演算自体は FP32 で実行されます
- layer の最終出力に signed intermediate 幅での wrap-around を適用します
- 同じ layer 内では出力 activation の再量子化は行いません

例えば `apply(model, QuantizeIntermidiate(ab=8, wb=8, ib=17), targets=[torch.nn.Linear], filters=[])` は、signed 8-bit の入力・重みを使い、最終的な intermediate 出力に signed 17-bit の wrap-around を適用します。
