DistillerでDeepLearningのモデルを軽量化: Gradual Pruning編
DeepLearningのモデル軽量化の気になっていたライブラリを使ってみました。今回はざっくりと導入の仕方と簡単な使い方、ライブラリの仕組みなどを調べた内容を書きたいと思います。はじめて使う人のガイドになればと思います。
IntelのNeural Network Distiller。pruningや8-bit quantizationなど軽量化アルゴリズムのフレームワーク。PyTorchのモデルを軽量化してONNX出力。TensorBoardと連携したモニタリングもできて使い勝手良さそう。https://t.co/5qm2IoUuuA
— tkato (@_tkato_) 2018年5月1日
Distillerとは
PyTorch向けのモデル圧縮ライブラリです。以下のような特徴があります。
- 数種類の枝刈り(pruning), 量子化(quantization), 正則化(regularization)アルゴリズムを実装
- 既存の学習スクリプトのtraining loopに追加するだけで使える
- 設定はYAML。モデルのレイヤー単位でpruningのパラメータを変えるなど柔軟な設定。
- TensorBoardと連携した、モデルのweightや精度の可視化
実用的でよさそう。これをインストールして、サンプルを動かしていきます。
インストール
詳しくはREADMEに書いてあるけれど、GPUささったUbuntu使うのが一番簡単。環境はPyTorch公式のdockerイメージ(pytorch/pytorch/0.4_cuda9_cudnn7)をつかいました。
コンテナの中で以下を実行してインストール完了
$ git clone https://github.com/NervanaSystems/distiller.git $ cd distiller $ pip install -r requirements.txt
動作確認
コンソールで動かせるサンプル（distiller/examples）と、Jupyter Notebook(distiller/jupyter)が用意されてます。まずはREADMEのGetting Startedに記載のある、training onlyのサンプルを動かしてみます。これは、軽量化云々はせず、ただ学習するだけですが、動けば環境構築は成功。また、ここで学習したモデルを今回のpruningの実験に使っていきます。
$ cd distiller/examples/classifier_compression $ python3 compress_classifier.py --arch simplenet_cifar ../../../data.cifar10 -p 30 -j=1 --lr=0.01
デフォルトでTensorBoard用のログも出力するので、以下でTensorBoardを起動して閲覧できます
$ cd distiller/examples/classifier_compression $ tensorboard --logdir='./logs'
図の最下段のsparsityは”モデル全体における値が0のweightの比率”で、これが大きいほどweightを削れた（=pruningできた）ことになります。例えば、sparsity=80なら、モデルの80%のWeightが0ということです。今回はただ学習しただけなのでsparsityは0です。
デフォルトの設定での学習終了時の精度は、以下となります。とりあえずこれをベースラインにします。
Saving checkpoint --- test --------------------- 10000 samples (256 per mini-batch) Test: [ 30/ 39] Loss 0.824917 Top1 70.820312 Top5 97.916667 ==> Top1: 70.990 Top5: 97.930 Loss: 0.82
学習済みモデルは移動しておきましょう。
$ mkdir simplenet_cifar $ mv best.pth.tar checkpoint.pth.tar simplenet_cifar/
このcompress_classifier.pyはclassificationモデルの学習コードにDistillerの呼び出しを追加した汎用的なサンプルです。”こんな感じでdistillerをご自身のtrain loopに組み込んで使ってね”という感じで、これを読めばなんとなくdistillerが理解できます。
ヘルプを見ると、以下のようになっており、いろいろ設定を変えられることがわかります。
$ python compress_classifier.py -h
usage: compress_classifier.py [-h] [--arch ARCH] [-j N] [--epochs N] [-b N]
[--lr LR] [--momentum M] [--weight-decay W]
[--print-freq N] [--resume PATH] [-e]
[--pretrained] [--act-stats] [--param-hist]
[--summary {sparsity,compute,optimizer,model,modules,png}]
[--compress [COMPRESS]]
[--sense {element,filter}] [--extras EXTRAS]
[--deterministic] [--quantize] [--gpus DEV_ID]
[--name NAME]
DIR
Distiller image classification model compression
positional arguments:
DIR path to dataset
optional arguments:
-h, --help show this help message and exit
--arch ARCH, -a ARCH model architecture: alexnet | densenet121 |
densenet161 | densenet169 | densenet201 | inception_v3
| mobilenet | mobilenet_025 | mobilenet_050 |
mobilenet_075 | resnet101 | resnet152 | resnet18 |
resnet20_cifar | resnet32_cifar | resnet34 |
resnet44_cifar | resnet50 | resnet56_cifar |
simplenet_cifar | squeezenet1_0 | squeezenet1_1 |
vgg11 | vgg11_bn | vgg13 | vgg13_bn | vgg16 | vgg16_bn
| vgg19 | vgg19_bn (default: resnet18)
-j N, --workers N number of data loading workers (default: 4)
--epochs N number of total epochs to run
-b N, --batch-size N mini-batch size (default: 256)
--lr LR, --learning-rate LR
initial learning rate
--momentum M momentum
--weight-decay W, --wd W
weight decay (default: 1e-4)
--print-freq N, -p N print frequency (default: 10)
--resume PATH path to latest checkpoint (default: none)
-e, --evaluate evaluate model on validation set
--pretrained use pre-trained model
--act-stats collect activation statistics (WARNING: this slows
down training)
--param-hist log the paramter tensors histograms to file (WARNING:
this can use significant disk space)
--summary {sparsity,compute,optimizer,model,modules,png}
print a summary of the model, and exit - options:
sparsity | compute | optimizer | model | modules | png
--compress [COMPRESS]
configuration file for pruning the model (default is
to use hard-coded schedule)
--sense {element,filter}
test the sensitivity of layers to pruning
--extras EXTRAS file with extra configuration information
--deterministic, --det
Ensure deterministic execution for re-producible
results.
--quantize Apply 8-bit quantization to model before evaluation
--gpus DEV_ID Comma-separated list of GPU device IDs to be used
(default is to use all available devices)
--name NAME, -n NAME Experiment nameまた、各軽量化アルゴリズムの設定は--compressでYAMLファイルで与える思想となっています。これを書き換えることでPruningのアルゴリズムを変えたり、別のPruningや量子化と組み合わせたり、レイヤ単位でのpruningの設定や、何epochからpruningを始めるか、などかなり細かい設定をファイルベースで行えます。実験の再現もしやすいですね。素敵。
Gradual Pruningのサンプルを動かしてみる
つづいて、実践的なpruningのサンプルを動かすことにします。
そもそもpruningとは、最終的な計算結果に寄与しないWeightを除去する手法で、「寄与」の指標としては単純にWeightの絶対値が使われたりします（けど研究によっていろいろ）。
Weightが減ると何が嬉しいのかというと、モデルのサイズや実行時のメモリ使用量を減らすことができ（フレームワークの実装依存でもあるけれど）、計算量削減や高速化ができる場合があります。
今回試すのは、Gradual Pruningというアルゴリズムで、1-stageのpruningです。よくあるpruningだと、学習済みモデルをpruningしてから再学習をする2つのstageを繰り返しますが、この方法では学習しながら絶対値の小さいWeightを0にする（マスクする）ことで、1回の学習の中でpruningまでしてしまうもの。
Gradual（=段階的な）の名前の通り、徐々にマスクをかけていくのですが、最終的なsparsity(Weightにおけるゼロの比率)の目標値を与えることができます。また、何epochかけて目標のsparsityへ到達するかもハイパーパラメータとして与えます。
Gradual Pruningの特徴としては、学習のステップ毎sparsityを以下のような関数として定義していることです。
これは、Weightが多く冗長なものも多い学習初期はスピーディに削り、Weightが少なくなってからは少しずつ削るような関数になっています*1。
ちなみにこれは以下のnotebookの実行結果で、ハイパーパラメータによって この関数がどう変化するかを確かめることができました。
https://github.com/NervanaSystems/distiller/blob/master/jupyter/agp_schedule.ipynb
pruning直後は少なからず精度が落ちるので、「pruningで精度が落ちる→再学習で回復」をちょっとずつ、良い感じのバランスで繰り返すことが重要になるのだと思います。論文によれば100~1000stepくらいごとにpruningするのが良いとのこと*2。
Distillerの公式ドキュメントには、Gradual Pruningを始め、各アルゴリズムの説明が結構詳しく書いてあります。勉強になった。
https://nervanasystems.github.io/distiller/algo_pruning/index.html#automated-gradual-pruner-agp
さて、脱線しましたがサンプルを動かしていきます。
先ほどのサンプルで学習したsimplenet_cifarモデルをGradual Pruningで軽量化してみます。先ほどのcompress_classifier.pyから実行できます。
YAMLファイルは、/examples/agp-pruning/以下などのサンプルをベースにして、以下のように記述しました。
simplenet_cifarは、conv1-conv2-fc1-fc2-fc3 という簡単なモデルです。今回は、各層の目標のsparsityを30%-50%-80%-80%-80%になるようにしています。pruningは、削る層によっては大幅な精度劣化を引き起こす*3ので、このように層毎にsparsityを変えられるのは便利です。また、100epochかけて徐々にpruningする設定です。
Gradual Pruningに限らず、YAMLでどんな設定がかけるのかについては、以下のドキュメントがわかりやすいです。
Compression scheduling - Neural Network Distiller
YAMLがかけたので、先ほど保存した最も精度が良いモデルに対して、Gradual Pruningを実行します。とりあえず200epoch。
$ cd distiller/examples/classifier_compression $ vim simplenet_cifar.schedule_agp.yaml $ time python3 compress_classifier.py --arch simplenet_cifar ../../../data.cifar10 -p 50 --lr=0.001 --epochs=200 --resume=simplenet_cifar/best.pth.tar --compress=simplenet_cifar.schedule_agp.yaml
学習完了後のTensorBoardはこんな感じ。青がpretrainで、ピンクがGradualPruning時のplotです。pruningによってsparsityが大きくなり、学習開始時は少し精度が低くなりましたが その後の学習で回復し、最終的にはpruning前と同等の精度に戻っています*4。
Saving checkpoint --- test --------------------- 10000 samples (256 per mini-batch) ==> Top1: 71.330 Top5: 97.880 Loss: 0.824
層毎のsparsityを見ると、確かに30%-50%-80%-80%-80%になってますね。
最後に、この軽量化できたモデルの情報を詳しく見てみましょう。compress_classifier.pyに対して、--summaryオプションをつけるとモデルの情報を表で出力してくれます。--summary=sparsityはそのままの意味ですが、--summay=computeとすると、層毎、全体の計算量を出力します。
Before
$ python3 compress_classifier.py --resume=./best.pth.tar -a=simplenet_cifar ../../../data.cifar10 --summary=sparsity $ python3 compress_classifier.py --resume=./best.pth.tar -a=simplenet_cifar ../../../data.cifar10 --summary=compute
After
$ python3 compress_classifier.py --resume=./checkpoint.pth.tar -a=simplenet_cifar ../../../data.cifar10 --summary=sparsity $ python3 compress_classifier.py --resume=./checkpoint.pth.tar -a=simplenet_cifar ../../../data.cifar10 --summary=compute
最後の行の"Total sparsity"に注目すると、pruningによって78.47%のweightを削減できたことがわかります。しかし、before/afterで計算量が変わっていないのは、weightは0であるものの行列自体のサイズは変わっていないからですね。値が0のweightを行列から削除して行列を組み替えることで、計算量自体を削減できます*5。このweightを実際に小さくする機能は、制限はあるもののDistillerに入っているようです。気になる方は”Model thinning”でDistiller内を検索してみてください
Distillerの仕組み
最後に、Distillerの仕組みについて調べた内容です。
Distillerの使い方は、設定をYAMLで書いてSchedulerをtraining loopで呼び出すだけ。
内部はどのようになっているのでしょうか。
ドキュメントだとこの辺に書いてあります。
https://nervanasystems.github.io/distiller/design/index.html
https://nervanasystems.github.io/distiller/schedule/index.html
仕組みは単純で、training loopの各イベントの前後でSchedulerのAPIを呼んでいます。
For each epoch:
compression_scheduler.on_epoch_begin(epoch)
train()
validate()
save_checkpoint()
compression_scheduler.on_epoch_end(epoch)
train():
For each training step:
compression_scheduler.on_minibatch_begin(epoch)
output = model(input_var)
loss = criterion(output, target_var)
compression_scheduler.before_backward_pass(epoch)
loss.backward()
optimizer.step()
compression_scheduler.on_minibatch_end(epoch)
上記の各on_xxxのタイミングで、SchedulerはPolicy(Pruning, Quantization, Regularizationを抽象化したクラス)を実行していきます。
Policyは、YAMLから設定を読み込んだ際に自動で登録されます。もちろんコード中で手動登録もできる。
PolicyにもSchedulerと同名のon_xxxメソッドが実装されており、SchedulerがコールされるとSchedulerに登録したPolicyもコールされます。Policyの種類により、schedulerのどのコールバックに対応するかが決まっています(反応しないやつはpassする)。
例えばPruningなら、on_epoch_beginでpruningするためのマスクを作成し、on_minibatch_beginではそのマスクをweightに適用し、weightの一部を0にします。マスク自体の作成の仕方は各Pruningのアルゴリズム(Prunerクラス)によって異なり、そこは分離されています。今回はGradual Pruningを使いました。
このように、種類が異なるモデル圧縮のアルゴリズムを共通の仕組みで扱える用になっており、拡張が容易ですね。
おわりに
再度ポイントをまとめて、締めたいと思います。
- compression schedulingの柔軟性
- レイヤーによって、どれくらいweightを削減しても精度劣化がないかは変わるので、レイヤー単位で細かく設定できるのは便利です。
- 新しい軽量化アルゴリズムを追加し易い設計
- training loopの各イベントに反応するPolicyとして抽象化しているのはわかりやすい。
- TensorBoardなどによるモニタリング機能の充実
- lossや精度だけでなく、実際にweightの統計情報をモニタできるのは便利です
- 過去の実験結果とも比較しやすいのもいいですね
(記事のタイトルにGradual Pruning編、と書きましたが次は未定です..)
