モデルのプルーニング手法を実装する

PyTorch Lightning によるスケーラブルな AI モデルの構築

Sergiy Tkachuk

Director, GenAI Productivity

プルーニングを使うタイミング

$$

📱 エッジデバイスや組み込みシステムへのモデルのデプロイに役立ちます

$$

➕ 量子化と組み合わせることで、効率をさらに高められます

$$

⚡ レイテンシやモデルサイズの削減を優先する場合に使用します

モデルのプルーニングとは

$$

ニューラルネットワーク内の重要度が低い接続を削除します
保存や計算が効率的なスパースモデルが得られます
一般的な手法として L1 unstructured pruning があります

プルーニングの例

モデルのプルーニングとは

import torch.nn.utils.prune as prune

prune.l1_unstructured(model.fc, name="weight",
                      amount=0.4)

print(model.fc.weight.data)

tensor([[ 0.25, -0.13,  0.05,  0.70],
        [-0.88,  0.31, -0.02,  0.44]]) # Before pruning


tensor([[ 0.25, -0.13,  0.00,  0.70],
        [ 0.00,  0.31,  0.00,  0.44]]) # After pruning (40% weights set to 0)

プルーニングマスクを理解する

$$

プルーニングは、対象となる重みテンソルごとにバイナリマスクを追加します。

$$

Mask = 1 --> 重みを保持します。
Mask = 0 --> フォワード時に重みを 0 に設定します。

$$

マスクを削除するまで、重みはメモリに保持されたままです。

プルーニングを確定する

デフォルトでは、プルーニングされた重みは元のテンソルの一部として残ります
プルーニングを確定するには、再パラメータ化を削除します
スパースレイヤーを、重みがゼロ化された標準レイヤーに変換します

Sequential(
  (fc): Linear(
    in_features=128, out_features=64,
    bias=True
    (weight): PrunedParam()
  )
) # Before prune.remove

import torch.nn.utils.prune as prune

prune.remove(model.fc, 'weight')

# Print model structure
print(model)

Sequential(
  (fc): Linear(in_features=128,
               out_features=64,
               bias=True)
) # After prune.remove

プルーニングの影響を評価する

元のモデルとプルーニング後のモデルの性能を比較します
精度はわずかに低下しますが、サイズやメモリは大幅に削減されます
このトレードオフがデプロイに適しているかどうかを判断する材料になります

プルーニングモデルのトレードオフ

練習しましょう！

PyTorch Lightning によるスケーラブルな AI モデルの構築

Preparing Video For Download...