TensorFlowをビルドしてみた

こんばんは、12/26以降の年末ブログです。[Champagne]時代からのファンです。

アドベントカレンダー終わったのにまだブログ書いてるのかって思いましたか?

弊社ではいま、ある目標のために有志が集まってブログを書いています。
これもそのうちの一本です。
大丈夫、まだ2021年だ。

今年はなんだかたくさん記事を書いたなあなんて振り返りつつ、特にエモい話をする予定はありません。
今日はTensorFlowをわざわざ自分でビルドしてみたのでその話でもしようかなと。
ところで皆さんTF派?PT派?(ミスター派?洋さん派?のノリでお願いします)

TensorFlowをソースからビルドする

いままでTensorFlowしか使ってなかったんですが、最近はもっぱらPyTorchを使うようになりました。
どうもCTO室AI推進部@ken11です。
PyTorchを使う機会が増えたのはひとえに僕が依存しているhuggingfaceがその方が楽だからですね。
時と場合によってもちろんTensorFlowも使っています。

さて、TensorFlowをソースからビルドする方法ですが、公式ではこちらで解説されています。
このリンク読めばいいだけ、といえばそれまでなんですが、せっかくなので少し紹介したいと思います。

※Linux環境かつGPU無しが前提です

Bazel

依存関係をインストールするところはそのままドキュメント通りなんでわかりやすいんですが、いきなり「Bazel使うよ」って言われるんですよね。
お恥ずかしい限りですが僕はこの「Bazel」を存じてなかったので、なにを入れたらいいのかわからず戸惑いました。
ドキュメントには「Bazeliskがあれば簡単に使えるよ」って書かれてるんですが、そのBazeliskのインストール方法までは書かれていないし僕の頭の中はもうバゼルギウスがハイパーバジリスクタイムでした。

Bazeliskはnpmでインストールできるらしいので、npm iする気でいたらそもそもnpm(というかnode)が入ってないということまで判明して、最終的に以下のようにしました。
いまだともう16でよかったのかな。

curl -sL https://deb.nodesource.com/setup_14.x -o nodesource_setup.sh | bash -
npm install -g @bazel/bazelisk

これで bazelisk コマンドが使えていれば問題ないようです。

Bazel build configs

その後、 ./configure まで終わらせると、いざ bazel build でビルドするところになるのですが、この際にオプションをいろいろ指定できます。
自分でTensorFlowをビルドする醍醐味はここだと思うのですが、 .bazelrc というファイルに指定可能なオプションの一覧があります。
しっかり確認して指定するとよいと思います。
ちなみに、今回は特にオプションを指定せずビルドしたものと、一部指定したものを作ってみました。

なお、 .bazelrc は最初からtensorflowリポジトリに存在するので、GitHub上でも確認することができます。

いざビルドしてみる

出オチみたいなこと言いますが、貧弱なマシンパワーでは普通にコケるし、成功しても3時間とかかかります。
いろいろ覚悟して bazel build コマンドを叩く必要があります。
ちなみに僕の経験ではメモリが8GBではコケました
ドキュメントに書かれてますが、メモリが足りない場合は --local_ram_resources=2048 を指定すると幸せになれるようです(未検証)

以下の画像はビルド結果の一例として、AWSの r5.large インスタンス(メモリ16GB)ではうまくビルドできたものの、3時間超かかりました。

ビルドが完了したらその後もドキュメント通りにpipで使える .whl にして終わりです。
これでようやくTensorFlowをインストールする準備ができました。

ビルドしたTensorFlowを使ってみる

pipでインストールして使ってみましょう。

pip install tensorflow-2.7.0-cp37-cp37m-linux_x86_64.whl

これでいつも通りに使ってみます。
早くなったりするんかな?というのも含めて比較してみたので紹介します。
比較はBERTのモデルである文章のFill-Maskタスクを10回ぶん回すという形でやっています。

  • 普通にpip installで降ってくるもの
  • オプションなしでビルドしたもの
  • avx2_linux オプションを指定してビルドしたもの

手持ち~ポケモン~TensorFlowはこの3種です、対戦よろしくお願いします
ファイッ!!

普通にpip installで降ってくるもの

これはなにも考えずに

pip install tensorflow

でインストールしたものですね

パッケージサイズは489.6 MBとなかなか大きいです。
サポートされてるものが多いからですかね(いまのバージョンって標準でもCUDAサポートとか入ってるので)

This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA

一昔前はAVXとかの命令は自分でビルドしないと有効にならなかった気がするんですが、いまは普通にpip installしても有効なんですね。

推論10回の実行結果は平均して0.53秒でした。

オプションなしでビルドしたもの

特にオプションを指定せずビルドしてみたものです。
パッケージサイズは159MBとなりました。先ほど比べると小さい!
今回はGPU無しでCPU推論しか考えてないので、CUDAサポートとかない分小さくなってると思われます。
これ、CPU推論で使う場合は明確にメリットとなりますね。
これだけでコンテナサイズ減らせると思いました。

This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA

てっきりオプション付けなければいずれの命令にも対応しないのかと思いきやなんかたくさん付いてきました。

推論10回の平均は0.55秒と、先ほどより遅くなってます。
対応の命令セットが増えた結果、適切な命令でない部分が出てきてしまったんでしょうか?
さすがにこの辺は僕にはわからないですが、とりあえず自分でビルドしたからといってデチューンになる可能性もあるというのがよくわかりました。

avx2_linux オプションを指定してビルドしたもの

最後に、先ほどのオプションなしのパターンをふまえて「このオプション付けておけばいいんじゃない…?」と思った avx2_linux を付けてビルドしたパターンです。

パッケージサイズは161MBでオプションなしとほぼ変わりませんでした。
コンテナサイズ減らすならこれでも十分ですね。

This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX512F FMA

avx2_linux だから AVX2 AVX512F FMA になるかなと思ってたんですが全然そんなことなかったです。
でもAVX512はAVX2の後継だと思うので、正直これで十分なんじゃないでしょうか。

推論10回の平均は約0.52秒、ここまでで最速の結果となりました。
微々たる差なのは~僕が横着して~タスクが適当なのが悪いと思っていて、タスク次第でもっと効いてくるのではないでしょうか。

まとめ

TensorFlowを自分でビルドしてみてわかったのは、とにかく時間がかかるということです←
もちろんマシンパワーがあればもっと早いんでしょうが…

まあ冗談は置いておいて、自分でビルドすることで、不要なサポートを切ってサイズをコンパクトにできたり、不要な命令を削って若干速度を上げることができるようだというのがわかりました。
ちなみにEC2Instances.infoというサイトを見ると、AWSの各インスタンスタイプでAVX2やAVX512がサポートされているかどうかを確認することができるのでとても便利でした。

逆にデフォルトだと切られているS3のサポート(noaws=true)等を変更してサポート付けるとなにができるようになるのか、僕はまだ知らないのでまた今度使ってみたいなと思いました。

そんな感じで、AI推進部では2022年もAI/MLの活用を推進していきます!
今年も大変お世話になりました。
みなさんよいお年をお迎えください。


マネーフォワードでは、エンジニアを募集しています。
ご応募お待ちしています。

【サイトのご案内】
マネーフォワード採用サイト
Wantedly
京都開発拠点

【プロダクトのご紹介】
お金の見える化サービス 『マネーフォワード ME』 iPhone,iPad Android

ビジネス向けバックオフィス向け業務効率化ソリューション 『マネーフォワード クラウド』

おつり貯金アプリ 『しらたま』

お金の悩みを無料で相談 『マネーフォワード お金の相談』

だれでも貯まって増える お金の体質改善サービス 『マネーフォワード おかねせんせい』

金融商品の比較・申し込みサイト 『Money Forward Mall』

くらしの経済メディア 『MONEY PLUS』

Pocket