rasnatune 0.0.3

Runtime compression tuning utilities for PyTorch

rasnatune

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

インストール

pip install rasnatune

クイックスタート

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 ごとに内部で複製されます。

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 を除外します。

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 を適用します。