rasna 0.0.1

Rasna: Realtime Adaptive Sparse Neural Accelerator

Rasna Python

Rasna は PyTorch モデルを Rasna runtime で実行するための Python ライブラリです。主な入口は次の 4 つです。

• rasna.tune(): Rasna 向けの疎化と量子化の影響を入れて、モデルをファインチューニングできる状態にします。
• rasna.compile(): eval() 済みの PyTorch モデルと sample input から Program を生成します。
• Program.save() / rasna.load(): コンパイル済み program を .rasna ファイルとして保存、再読み込みします。
• Runtime / RuntimeCSim: 別途配布される runtime object と Program を組み合わせて推論を実行します。

このディレクトリ以下は PyPI package rasna として配布されます。bitstream や simulator shared library などの runtime object は Python package には含まれず、別途配布されます。

License

Rasna は proprietary software です。オープンソースライセンスではありません。商用利用、改変、再配布、リバースエンジニアリングは禁止されています。詳細は同梱の LICENSE を確認してください。

Installation

pip install rasna

Rasna は NumPy と PyTorch に依存します。CUDA 版など特定の PyTorch wheel が必要な場合は、先に環境に合った PyTorch を入れてから Rasna を入れてください。

pip install torch torchvision
pip install rasna

基本フロー

Rasna を初めて使うときの流れは次の通りです。

1. PyTorch モデルを用意する、または checkpoint から読み込む。
2. rasna.tune() を呼び、Rasna 向けにファインチューニングできる状態にする。
3. 通常の PyTorch training loop でファインチューニングする。
4. model.eval() に切り替える。
5. 代表的な batch-first sample input と一緒に rasna.compile() へ渡す。
6. 生成された Program を .rasna ファイルに保存する。
7. 後で .rasna と runtime object を読み込み、runtime で実行する。

rasna.tune

rasna.tune() は supported layer を in-place に変更し、Rasna の実行条件に近い量子化効果と sparsity mask を入れた状態でファインチューニングできるようにします。

import torch
import rasna

model = MyModel()
model.load_state_dict(torch.load("checkpoint.pt", map_location="cpu"))

rasna.tune(
    model,
    sp=0.8,
    ab=8,
    wb=8,
    ib=17,
    chpad=True,
    strategy="auto",
)

model.train()
for images, labels in train_loader:
    optimizer.zero_grad()
    loss = criterion(model(images), labels)
    loss.backward()
    optimizer.step()

主な引数は次の通りです。

• sp: weight pruning の sparsity ratio。既定値は 0.8 です。
• ab: activation quantization bit width。既定値は 8 です。
• wb: weight quantization bit width。既定値は 8 です。
• ib: tuning 中の中間値 bit width。既定値は 17 です。
• chpad: Conv2d の padding mode を circular にします。既定値は True です。
• strategy: "auto" は first convolution、depthwise convolution、linear layer の sparsity を避けます。"all" は supported Conv2d / Linear すべてに sparsity を適用します。

ファインチューニング後の PyTorch checkpoint は、通常の PyTorch モデルと同じように保存できます。

rasna.compile

rasna.compile() は eval() 済みの PyTorch モデルと、1 つ以上の batch-first sample input を受け取り、Program を返します。この関数はモデルを実行する runtime object を生成しません。

import torch
import rasna

model.eval()

# 4D input は batch-first です。既定では channel-last の BHWC を使います。
sample = torch.randn(16, 32, 32, 3)

program = rasna.compile(model, sample)
program.save("model.rasna")

PyTorch model を channel-first BCHW で書いている場合は、channel_last=False を指定します。

sample = torch.randn(16, 3, 32, 32)
program = rasna.compile(model, sample, channel_last=False)

複数入力モデルでは、sample input を位置引数としてすべて渡します。

program = rasna.compile(model, sample_left, sample_right)

現在 trace できる主な演算は、Conv2d, Linear, ReLU, BatchNorm2d, MaxPool2d, AdaptiveAvgPool2d, torch.flatten, tensor add/sub, scalar affine operation です。未対応の演算がある場合は、compile 時に FX node 名つきで報告されます。

Program の save / load

Program は Rasna のコンパイル済み artifact です。入力・出力の説明と、runtime に渡す instruction stream を保持します。

program.save("model.rasna")

loaded = rasna.load("model.rasna")

program file には .rasna 拡張子を使います。コンパイルに時間がかかる場合、training code と deployment code を分けたい場合、または別のマシンで runtime 実行したい場合は、Program.save() で保存しておくと便利です。

Runtime

runtime 実行には、コンパイル済み Program と、その program に対応する runtime object が必要です。Python package には runtime object は含まれません。

• simulator runtime object は通常 .so shared library として配布されます。
• device runtime object は通常 .bit bitstream と、その実行に必要な付随ファイルとして配布されます。

Program から runtime を作ります。

import numpy as np
import rasna

program = rasna.load("model.rasna")
runtime = program.to_runtime("/path/to/runtime-object.so")  # simulator
# runtime = program.to_runtime("/path/to/runtime-object.bit")  # device

x = np.random.randn(2, 32, 32, 3).astype(np.float32)
y = runtime(x)

program.to_runtime() は .so path では RuntimeCSim を、それ以外の runtime object path では Runtime を生成します。bitstream を使う device 実行では、配布された runtime object を読み込める Python 環境、たとえば PYNQ 環境が必要です。

runtime input は leading batch dimension つきでも渡せます。compile と実行を channel-first でそろえる場合は channel_last=False を指定します。

y = runtime(x_nchw, channel_last=False)

既定では、runtime は compile 時に決まった fixed-point type に合わせて入出力を normalize します。すでに packed 済みの int8 / uint8 値を渡す場合は normalize=False を使います。

y_int = runtime(x_int8, normalize=False)

End-to-End Example

import torch
import rasna

model = MyModel()
model.load_state_dict(torch.load("checkpoint.pt", map_location="cpu"))

rasna.tune(model, sp=0.8)

model.train()
for images, labels in train_loader:
    optimizer.zero_grad()
    loss = criterion(model(images), labels)
    loss.backward()
    optimizer.step()

model.eval()
sample = torch.randn(8, 32, 32, 3)

program = rasna.compile(model, sample)
program.save("model.rasna")

runtime object がある環境では、保存済み program を読み込んで実行できます。

import numpy as np
import rasna

program = rasna.load("model.rasna")
runtime = program.to_runtime("/path/to/rasna-runtime.so")

inputs = np.random.randn(1, 32, 32, 3).astype(np.float32)
outputs = runtime(inputs)

モデル作成時の注意

• rasna.compile() の前に必ず model.eval() を呼んでください。
• rasna.compile() に渡す sample input には leading batch dimension が必要です。
• sample input の値は fixed-point type や buffer layout の決定に使われるため、実運用に近い代表値を使ってください。
• 未対応の PyTorch operation は compiled model の外に出すか、supported module / tensor operation に書き換えてください。
• model、input shape、target configuration、runtime object が変わった場合は再コンパイルしてください。