モデルの蒸留を実装し freesound2019 コンペで検証してみた。

こんばんは、kaerururu (@kaeru_nantoka)です。

今回は、Distillation the Knowledge in a Neural Network (2015) [ https://arxiv.org/pdf/1503.02531.pdf ] を読みました。

そして、kaggle freesound2019 コンペで実際に使ったデータとモデルを用いて蒸留の検証をしたので、これについて書いていこうと思います。

通しの実装は、kaggle 上に公開カーネルとして載せていますので合わせてご覧ください。

Distillation : [ https://www.kaggle.com/kaerunantoka/distillation-implementation-with-freesound2 ]

Simple (比較用) : [ https://www.kaggle.com/kaerunantoka/simple-cnn-image-size-299 ]

参考

実装の参考

https://github.com/moskomule/distillation.pytorch/blob/master/hinton/utils.py

温度付き softmax 関数の理解する際に参考

https://qiita.com/nkriskeeic/items/db3b4b5e835e63a7f243

KL divergence を理解する際に参考

https://qiita.com/ceptree/items/9a473b5163d5655420e8

概要

まず、蒸留の基本的なお気持ちについてご紹介します。

一般的に、教師データをこれ以上増やせないという制約の元でより良い精度を実現するためにはモデルのアンサンブルという手法が取られます。この手法はより良い予測を得られる一方、合計のモデルの容量がとても大きくなり、しばしばデプロイ環境 (モバイルデバイスなど) で利用するのに不便が生じます。

そこで、より良い予測を得られるアンサンブルモデルの出力を別な軽量なモデルの学習にうまいこと利用しようという発想です。

以下、

・教師モデル (アンサンブルモデルなど容量が重いがより良い予測が得られる。)

・生徒モデル (より軽量なモデル。教師モデルの出力も学習に使う。)

とします。

本記事の目標

これまでの説明を聞いて、「より良い予測が得られるが、モデルが重い... 」

スーパーアンサンブルモデルを送りつけられる kaggle の host を想起する読み手のかたも多くいらっしゃると思います。

そこで本記事では、モデルの蒸留を実際に筆者が取り組んだ kaggle freesound2019 コンペで使用し、金メダルをとったスーパーアンサンブルモデルから、本コンペでベースラインとなったシンプルな CNN モデルへと知識を継承させる実験をしました。

結果

実験の結果以下のようになりました。

	CV	Private LB	順位	合計容量
A	／	0.74435	10位 (金)	188MB (94 * 2)
B	0.75014	0.59023	190位	18MB
C	0.74545	0.60895	185位	18MB

model A : 教師モデル, InceptionV3 * 2 のアンサンブル

model B : 生徒モデル1 (素)

model C : 生徒モデル2 (教師モデルの出力使う)

LB スコア… 私の実験ではおよそ 0.018 の上昇で、順位だと5位アップ相当でした。教師モデルのスコアが 0.744 なことを考慮すると、イマイチパッとしないように見えますが後ほどこれに対する私見を述べます。

CV スコア … 少し下がっています。実はこれもまた蒸留の１つの特徴で、overfit に抑制がかかるとされています。言い換えると、蒸留によってある種の正則化の効果が得られます。これは当論文の第6章に書かれておりますのでご覧ください。

ファイルサイズ … 教師モデル Inception V3 はそれぞれ 94MB で合計 188MB でした、一方で生徒モデルの CNN は 18MB でした。

より強力なモデルの力を借りることで、軽量なモデルの予測が少し良くなっているのがみて取れます。

実装の説明

ポイントは In[23] の以下の２箇所です。

for x_batch, y_batch in progress_bar(train_loader, parent=mb):
            
    teacher1 = model1(x_batch[:, 0, :, :].view(-1, 1, 299, 299).cuda())
    teacher3 = model3(x_batch.cuda())
    lesson = ((teacher1 + teacher3) / 2)
    preds = model(x_batch.cuda())
            
    kl_loss = F.kl_div(F.log_softmax((preds / temperature)), F.softmax((lesson / temperature)),
                               reduction="batchmean")
    loss = criterion(preds, y_batch.cuda()) + lambda_factor * (temperature ** 2) * kl_loss

    model = Classifier(num_classes=num_classes).cuda()
    
    model1 = model_1ch_incep
    model1.load_state_dict(torch.load(f"../input/inception-1ch-dameoshi-aug/weight_best{binary_number}.pt"))
    model1.cuda()
    model1.eval()    
    
    model3 = model_3ch_incep
    model3.load_state_dict(torch.load(f"../input/inception-3ch-dameoshi-aug/weight_best{binary_number}.pt"))
    model3.cuda()
    model3.eval()

蒸留では、教師モデルと生徒モデルの出力の分布を似せるように学習させます。そのため、同じ iter の中で prediction をするために上のように一気に model と重みを load します。

そして得られた両方の出力は通常の softmax 関数ではなく温度付き softmax 関数 (softmax_with_temperature) に通します。

というのも蒸留で利用したい教師モデルが学習した情報というのは、出力値の分布 (何々と何々は間違えやすいといったもの)であるからです。通常の softmax 値だと正解ラベル以外の確率が限りなく０に近い値が割り振られてしまうため、正解ラベルを利用するのとさして変わらず不適当なので、温度 (=重み) をつけてあげて低い確率を重視しようぜというお気持ちなのです。

それらの'ならされた'出力たちを KL divergence Loss 関数に渡して、これを加味した CrosEntropyLoss を最小にしようという流れです。

KL divergence は２つの確率分布の類似具合を測る指標で、２つの確率分布が完全に一致したら (=lossがなくなったら) 0 になる特徴があります。

これを最小化しようというのはなるほど納得という感じです。

結論

教師モデルが学習した知見をモデルサイズのより小さい生徒モデルに受け継がせることで何もしない生徒モデルに比してより良いスコアを得る、という蒸留の望む効果を kaggle の過去コンペで実際に使用したアンサンブルモデルに対して適用しても一応の結果が得られました。

考察

教師モデルのスコア 0.74435 に対して、生徒モデルのスコアは 0.60895 と効果は限定的なのではないかと感じます。このことについて私は、教師モデルの精度に近づくように生徒モデルが学習しているというより、生徒モデル本来の伸び代の限界までしか精度が伸びないといった側面があるのではないかと考えました。例えば、教師モデル SE-ResNeXt で生徒モデル 3層のMLP で学習させたケースにおいて、生徒モデルのMLPが SE-ResNeXt に匹敵する表現力を得られる、というのは考えにくそうであるなという所感です。

実際当論文では、入力のパラメータを変えた同じモデル 10 個のアンサンブルを教師モデル、同じモデルのシングルモデルを生徒モデルとし、そのままのシングルモデルよりよい精度が得られたといった文脈で書かれていました。

まとめ

・ CV は下がる (ある種の正則化効果) が、LB は微増 (汎化性能の向上) した。

・軽量なモデルの予測精度向上への寄与に成功した。

・一方で、教師モデルに近い表現力を持った生徒モデルを利用しないケースでは精度は向上するものの効果は限定的であるという示唆が得られた。

以上3点、今回私が実験した限りではわかりました。

特に「CV が下がる (が成功している)」という点が 'Trust CV' を信条とする多くの kaggler にとって実際のコンペで利用するのに心理的な壁になりそうだと思いました。一方で、昨今の kaggle では限られた推論時間にスーパーアンサンブルモデルをいかに詰め込むかといったことに注力するシーンが散見されます。本実験では扱いませんでしたが、同じモデルの複数アンサンブルの予測に近いものを同一モデルのシングルモデルに出力させるような当論文の手法は、限られた時間を有効に活用しなければならない現代 kaggle においての１つの解足りうるのではないかと思います。

以上誤り等に気づかれましたら、優しくご指摘くださいますようよろしくお願いいたします。ありがとうございました。

追記

2019/07/24 「実装の KL divergence loss の計算をしているところで、なぜ softmax ではなく log_softmax を使用しているのか？」という質問をいただきました。結局質問者の方が、F.kl_div(input, target) 関数に渡す input は log 形式である必要がある旨を docs を調べて教えてくださいました。　https://pytorch.org/docs/stable/_modules/torch/nn/modules/loss.html ありがとうございました🙇‍♂️

かえるのプログラミングブログ

プログラミングでつまずいたところとその解決策などを書いていきます。

モデルの蒸留を実装し freesound2019 コンペで検証してみた。

目次

参考

概要

本記事の目標

結果

実装の説明

結論

考察

まとめ

追記