Ryuzのブログ

今日は、Networkの蒸留(Distillation)について考察してみたいと思います。

LUT-Networkに限らないのですが、その課題に学習時のGPUのメモリの不足と、学習時間の問題があります。

例えばMobileNetという軽量なネットワークがあります。これは畳み込み層をpointwiseとdepthwiseに分けて行うことで係数の数を大幅に削減しています。
軽量にすると学習時も計算量も減るように思ってしまいますが、実際には処理を分割したことで層が増えてしまい、逆にGPUの消費メモリが増え、学習時のメモリへの読み書きも増えてしまいます。

この延長で、超SparseであるLUT-Netは大規模な学習が非常に困難です。そこで、大規模な学習がしやすいネットワークで学習してから、メモリに収まる範囲で少しづつLUT-Networkに写し取っていくというというアイデアを考えています。
これを蒸留(Distillation)と呼ぶそうです。

現在、取り組もうとしている流れが以下の図です。

まず、一般的な畳み込み(Convotution)ネットワークの畳み込み層をバイナリ化する場合、ReLUなどのActivation層を2値化を行うBinary-Activation層に置き換えることになります。この時、2値化の手前で BatchNormalizationを行うのが効果的です。

BatchNormalizationは通常 Inferrence時には単なるスケーリングのみとなります。スケーリングのみであるため、2値化と組み合わせる場合、2値化の閾値自体をずらしてやればスケーリングも不要となります。

しかしながら、ここで前提として計算の中間結果が非バイナリであることが求められます。GPUを用いる場合、GPU内の内部演算が一時的に多ビットになることは特にデメリットはありません。しかしながらFPGA化する場合、そのままリソース規模に効いてくるのでこれは致命的です。
また、画素サイズ分のスケーリング係数が発生するので規模の点でも非常に厳しくなります。

そこで、LUT-Networkでは、当初より BatchNormalizationとBinary-Activationを畳み込み処理に取り込んでいます。これは工夫したというより、こうするしかLUT-Net は成立しないので、特に気にすることなく当初からこうしていたわけですが、LUT-Netの特徴の一つでもあります。
一般的な畳み込み(Convotution)は、効率化のためにGPUなどで一括で計算されますが、処理量やメモリを増やして良ければ、Loweringを行い、im2col と col2im の処理に挟まれる形でDenseAffine計算に分解することができます。

Loweringを行うと、学習時のメモリ消費や計算時間は増えますが、畳み込みの中にDenseAffine以外の処理が入れ込めます。

今回考えているのは、従来のネットワークに、LUT-Networkで行っているこれらの仕組みのみを適用して学習させ、学習後に蒸留して取り込もうというアイデアです。

蒸留は、今のところ

• 元のDenseAffineの重み順で結線すると効果的
• 元のDenseAffineの結線数が少ないと模倣しやすい(おそらくMobileNetは効果的)

などが、見えており、もうひと頑張りといったところです。

なお、今回は蒸留を目的に従来のネットワークを変形して学習させようとしていますが、おそらくこの変形は蒸留することなくそのままFPGAにしても効果的ではないかと考えています。

DenseAffine + BatchNormalization + Binary-Activation はInferrence時に XNOR-Net のような構造にできるので、そのままHLSなどで合成すれば効果的にFPGA化できる可能性があります。

先日の記事で、他のバイナリネットワークに対するLUT-Networkのリソース規模の大雑把な比較を行っていますが、その際に前提としてLUT-Networkに対して行っている上記のような技法を適用しているのを前提としていた部分がありました。

またXNOR-Netについてはかなり概算が入っていたので、実際にいくつか合成実験も行ってみました。
私自身XNOR-Netは試したことがないので理解不十分かもしれませんが、理解の範囲で実験しています。

思っていたよりXNORネットはFPGA化において高効率に思います。

一方で全結線という点ではNに対してO(N log N)で規模増大する可能性を有しているのと、従来のパーセプトロンモデル以上のことはできない(単段ではXORが解けない)など、LUT-Netに比べて非効率な部分はあるかと思いますが、如何せん、PyTorchなどの既存のプラットフォームで学習できる可能性がある点で、手が付けやすいのは事実かと思います。

LUT-Net が一般的なネットと遠いので、なかなか蒸留できる接点を作っていくのが難しいのですが、FPGA化時の高効率化のためのポイントを押さえていけば、その中間的なところも組み立てていけそうな気もしますし、それ自体もいろいろと面白い検討ができそうな気がしています。