コンピュータビジョンの高速化：IPU-PODスーパーコンピューティングシステムに合わせたEfficientNetのスケールアップ

Graphcore Researchが、革新的なコンピュータビジョンモデルEfficientNet-B4の学習をGraphcoreの最新のハイパースケールIPU-POD128およびIPU-POD256システム上で高速化し、2時間弱で収束させた方法をご紹介します。

EfficientNetとは？

EfficientNetファミリーのモデルは、比較的少ないパラメータとFLOPsで高いタスク性能を達成する、コンピュータビジョンの最先端を示すものです。しかし、レガシーなプロセッサアーキテクチャではEfficientNetの効率の良さを引き出せないため、実際に採用するには限界があります。

例えば、EfficientNetsのdepthwise畳み込みやグループ畳み込みはどちらも、通常の畳み込み演算と比較すると表現力が高く、計算効率が高い一方で、演算強度（データ移動に対する演算の割合）が低いという特徴があります。このタイプの演算は、メモリとプロセッサコア間のデータ転送量が多いため、GPUでは性能が低下します。それに比べてメモリセントリックアーキテクチャを採用するIPUでは、モデル全体とそのアクティベーションをチップ上に残すことができ、そのようなコストの高いデータの移動が軽減されます。

さらに、MIMD（複数命令、複数データ）パラダイムを採用するIPUでは、モデルや学習手順の多次元にわたるきめ細かな並列処理が可能です。これにより、小さな畳み込みを用い、ごく少数のデータサンプルを並列に処理しながら、高いスループットを実現できます。それに対し、SIMD（単一命令、複数データ）アーキテクチャを採用するGPUでは並列処理に限界があり、機械学習アルゴリズムの設計とハードウェア上で達成可能なスループットの間でトレードオフを強いられます。

このようなIPUハードウェアでの高速化により、さらに多くのイノベーターやAI専門家が実世界の大規模実装にEfficientNetモデルの高効率を活用できるようになります。

GitHubで公開されているGraphcoreのサンプルでは、EfficientNetファミリーのモデルなど、誰でもすぐにIPUで使える人気のモデルが幅広く提供されています。 今回のエクササイズでは、AI分野でEfficientNetの性能のベンチマーク評価によく利用されるEfficientNet-B4を検討します。

スループットの最適化

モデルの分散

データIO

この実行スキームを使用するPoplarは、数十万タイルに及ぶ演算を効率的に並列化できます。IOが制限されるアプリケーションでは、すべてのタイルでナイーブなBSPパラダイムを使用すると性能が低下してしまいます。その点においてPoplarは、一部のタイルを同期させることなくデータのストリーミングに特化させ、残りの計算タイルはBSPパラダイムで実行することで、この問題に対処しています。IOをオーバーラップさせることでIOのボトルネックを軽減しつつ、高度にスケーラブルな実行スキームも可能にしているのです。私たちは、前述のメソッドと組み合わせて、IPUあたり1472タイルのうち32タイルをIOオーバーラップに割り当て、最大3つのデータバッチをプリフェッチすることで、すべてのシステムスケールで素晴らしいスループットを達成できることを明らかにしました。

バッチサイズのスケーリング

IPU-POD128やIPU-POD256のような大規模なシステムで学習するには、大きなグローバルバッチサイズを使用することが有効です。そのようにすることで、データサンプルの処理を多数のレプリカで並列化しながら、各レプリカに十分な作業量を確保して、効率を維持することができます。 さらに、レプリカ間で勾配を減らすコストをならすために、レプリカ間で通信して重みづけを更新する前に、複数の前進パスと後退パスで勾配を局所的に蓄積する、いわゆる勾配蓄積（GA）を行います。したがってグローバルバッチサイズは、ローカルバッチサイズ、レプリカ数、および勾配蓄積数の積になります。

 

ベースラインには、『Making EfficientNet More Efficient（EfficientNetをより効率的にする）』を参考に768のグローバルバッチサイズを使用しています。ローカルバッチサイズを3にした場合、IPU-POD256では1回の学習イテレーションで2つのローカルバッチしか得られません。これではパイプラインの設定を十分に活用できず、コストのかかる重みづけの更新を頻繁に行う必要があります。しかし、単純にグローバルバッチサイズを大きくすると、一般化の性能低下を招きます。この現象は、多くの機械学習アプリケーション、特にコンピュータビジョンに共通するものです。したがって、すべてのIPU-PODシステムで高いスループットを維持できる大きさでありながら、優れた統計的有効性も同時に達成できる小さなグローバルバッチサイズを求めることになります。

 

よく知られていることですが、オプティマイザは、大きなバッチサイズを使った学習のロバスト性に大きく影響することがあります。 元々のEfficientNetの論文ではRMSPropというオプティマイザが使用されていましたが、私たちの調査では、バッチサイズを大きくしたときに統計的有効性を維持するのに苦労しました。Wongpanichらも同様の見解を示しており、代わりに、バッチサイズの大きなビジョンモデルの学習に有効であることが知られているオプティマイザ、Layer-wise Adaptive Rate Scaling（LARS） を使うことを提案しています。LARSは、重みづけと重みづけ更新の大きさが同程度になるように、学習率を層単位でスケーリングします。また私たちは、オプティマイザとともによく使われる多項式減衰学習率スケジュールをLARSと併用しています。

 

私たちは、バッチサイズ、学習率、ウォームアップエポック数、モーメンタム係数、および重みづけ減衰のハイパーパラメータスイープを行った結果、当初のEfficientNet-B4実装に比べて性能を一切低下させることなく、6144のグローバルバッチサイズを達成できることを明らかにしました。このバッチサイズでは、検討したすべてのIPU-PODにおいて勾配蓄積数を高いまま維持できます。

性能結果

柔軟なハイパーパラメータ設定を使うことで、様々なIPU-PODシステムに対して高いスループットでEfficientNet-B4を学習できるようになりました。元々のEfficientNetの論文と同様に、350エポック分の学習を行います。『Making EfficientNet More Efficient（EfficientNetをより効率的にする）』では、EfficientNetファミリーのモデルをより低い解像度で事前学習してから、本来の解像度で微調整できることが実証されていますが、今回は他の実装と比較できるように、元の解像度を維持しました。

すべての実験において、IPU-PODのサイズに応じてレプリカ数と勾配蓄積数をスケーリングすることで、同じ基本的な機械学習のハイパーパラメータ構成を使用しました。その結果、82.54 ± 0.13％の検証精度でいずれも収束しました。

上記の結果から、主要なGPUハードウェアでの結果と対比して、IPUベースのシステムにおいてEfficientNetを学習させることがいかに高速化されるかがわかります。イノベーターが新しいアイデアを素早く反復・テストするためには、モデルを高速で学習することが不可欠です。その点、IPUベースのシステムではEfficientNet-B4を2時間以内に学習できるので、イノベーションのスピードを数日から数時間に短縮できます。

 

この高速化により、多くのコンピュータビジョンアプリケーションは、レガシーなプロセッサアーキテクチャの制約を受けることなく、大規模な学習を大幅に高速化できることが実証されました。EfficientNetのような次世代コンピュータビジョンモデルを、GraphcoreのIPU-PODシステムのような最新のハードウェアで使用すれば、CTスキャン解析やビデオアップスケーリングから故障診断、保険金請求の確認まで、ビジョン関連の膨大なユースケースを高速化できるでしょう。

 

これらの構成は現在、GraphcoreのGitHubサンプルでお試しいただけます。（注：すべてのIPU-PODシステムの最新構成は、当社のPoplar SDK 2.4のリリースに伴い更新されます。）

 

この記事はTowards Data Scienceに掲載されたものです。

Towards Data Scienceの記事を読む

謝辞

この研究にご協力いただいたDominic Masters氏とCarlo Luschi氏に謝辞を述べるとともに、洞察を与えてくれたGraphcoreの同僚からの協力にも謝辞を述べます。