M1 Mac + Dockerの環境でTensorflowをビルドする

pip install tensorflowできない！

M1チップのMacではpipでTensorflowをインストールできません。AArch64に対応したパッケージはPyPIで公開されていないようです。Macに直接インストールする方法はいくつかあるみたいですが、環境は汚したくないのでDocker上で使いたいところ。少し調べたところ、ビルドするしかないみたいだったので、ビルドしてみました！
使えるようになるまで結構ハマるポイントがあったので、備忘録として残しておきます。

公式ドキュメントにビルド方法が書かれているが

https://www.tensorflow.org/install/source?hl=ja
ここにビルド方法が書かれています。

1. ビルド環境を用意してソースコードからビルド
2. ビルド環境が用意されたDockerイメージ上でビルド

上記の2通りの方法がありました。Docker環境でTensorflowを使いたいので、まずは2の方法でビルドを試みました。が、ビルドできませんでした。ビルドツールのBazelがx64向けのイメージだったようで、Bazelが使えず。。。

ということで、M1 Mac + Dockerの環境でTensorflowを使うにはビルド環境を一から準備してビルドしないと使えないようです。

いくつかハマりどころがあります！

基本的には公式ドキュメントの通りに進めていけばいいのですが

1. Bazelのビルドオプション
2. 依存パッケージのh5py
3. numpyのバージョン

ここでハマりました。
なお、Bazelは公式ドキュメントに「Use Bazelisk (recommended)」とあったので、npmを使ってインストールしました。

Bazelのビルドオプション

bazel build [--config=option] //tensorflow/tools/pip_package:build_pip_package

このコマンドでビルドしますが、メモリ8GのMac Book Airではオプション指定無しだとビルド中にメモリが足りなくなり完了しませんでした。公式ドキュメントには--local_ram_resources=2048でRAM使用量を制限してくださいとありますが、このオプションだけではまだメモリの問題で完了しません。最終的に

1. メモリを最大8GBまで使えるようにDockerDesktopの設定を変更
2. オプションは--local_ram_resources=2048 --jobs=1を指定

とすることでビルドできるようになりました。Dockerの設定は8GB未満でもいいかもしれませんが、最大にしておくといいと思います。（ビルド完了後は元に戻すことを推奨）
--jobsオプションは並列数を指定するオプションで、少しでもメモリ使用量を減らしたいので1にしています。2でもギリギリいけるかもしれませんが、メモリ８GBの環境では1が無難です。

依存パッケージのh5pyも・・・

Bazelのビルドが完了すれば、wheelの作成までは問題なく進められると思います。しかし、作成したパッケージをpipでインストールするときに、またハマりました。依存パッケージのh5pyが、これも、AArch64に対応したパッケージがPyPIで公開されておらずpipでインストールできず、Tensorflowもインストールできません。調べてみると、幸いapt-getでlibhdf5-devをインストールすることで、h5pyが使えるようになるみたいだったので、その方法を取りました。

apt-get install -y libhdf5-dev

でインストール後、

pip install --no-binary=h5py h5py

とすることで、インストールできました。

参考: https://github.com/h5py/h5py/issues/1181

numpyのバージョン

これでようやくTensorflowをpipでインストールできるようになりました。
しかし！
import tensorflowでエラーが出ました...。
RuntimeError: module compiled against API versionだそうです。どうやらTensorflowが要求するバージョンのnumpyがインストールされた環境でビルドしないといけなかったようです。
今回はTensorflow2.6で、numpyは1.19を要求されているので、

pip install numpy==1.19.2

でnumpyをインストールし直して、再度Tensorflowをビルドしました。Bazelのビルドからやり直しましたが、wheelの作成し直しだけでいいかもしれません。

まとめ

AArch64環境でTensorflowをビルドする場合は以下の事項に注意！

1. ビルド環境に合わせてBazelのビルドオプションを指定する（--local_ram_resourcesだけでなく--jobsも使う）
2. numpyのバージョンはTensorflowのバージョンに合わせたものをインストールしないといけない
3. h5pyはapt-getでインストールした上で、pipの--no-binaryでインストール

ビルドできてしまえば大したことないのですが、結構大変でした。（Bazelのビルド時間が結構長い。。。）

ここまで説明しておいてなんですが、おそらくDockerHubを探せばAArch64環境用のTensorflowイメージがあると思うので、それを使うのが手っ取り早いと思います。
ただ、ビルド設定で色々カスタマイズしたパッケージができます。（普通に使う分にはデフォルト設定で十分ですが）

Tensorflowをソースコードからビルドしようとしている人の一助となれば幸いです。