MacでもDeep Learningがしたい。PythonからMetal APIを呼び出す簡易実験の話
【2023年 8月6日追記】
この記事は筆者が当時使っていたX86アーキテクチャで動作するMacに向けた記事であり、すでにレガシーなものとなっています。Apple silicon搭載のMacを利用しているPytorchユーザーの方は、PytorchのMPSバックエンドの利用を推奨します。
---------------------------- 追記終わり ----------------------------
これまで数学の話が多かったですが、今回は少しハッキーなプログラミングの話です。直近の仕事でGANという敵対性生成ネットワークを用いた画像生成を行うというものに少し触れたのですが、そのプロジェクトにおいてはGANのGeneratorに対する学習がクラウド上のGPUマシンでも1週間強かかるということがありました。これまで自然言語モデルのファインチューニング等はよく行っていましたが、それでも一回の学習で1週間かかったということは無く目標の精度を達成できていたので、改めてConvolutionの計算量の多さに驚かされたというものです。
現状、僕の手元にあるPCはMacのみでして、しかもAMDのRadeon搭載かつCUDAをサポートしていないという現状においては、Tensorflowなどの性能を満足に活かせないため、仕事においてもローカルでさくっと実験することが難しかったり、あるいは研究においても自腹でコストをかけてクラウドマシンを利用する必要があります。このことでNVIDIA製のGPU搭載マシンを買うことを検討していたのですが、まずはMacでのGPUプログラミングの可能性を探ってみることにしました。
MacでGPUプログラミング
冒頭でも述べたとおり、最新のMacOSXではCUDAをサポートしておりませんので、Macにおいてメジャーなディープラーニングフレームワークを利用する場合、もろもろの計算は通常はCPUで実行されることになります。そこでOpenCLを用いたGPUプログラミングも考えられるのですが、MacではMetalという、AppleのPlatformに最適化されたGPUプログラミングのためのAPIが存在します。実際、Metalには体系化されたCNN用のAPIが存在しておりまして、SwiftかあるいはObjective-C利用することで簡単にCNNの訓練と予測を行うことが可能です。（今回はSwiftを利用します）
SwiftでDeep Learningを行うということ
Swiftは大変良い言語で、アプリケーションを作ることに関しては本当に素晴らしいのですが、いわゆるデータ分析タスクのエコシステムという面においてPythonの足元にも及びません。（それは例えばpandasやnumpy、matplotlib等にきちんと相当するものがないということです）
そこで、PythonからSwiftを介してMetalにアクセスしようというのが今回の実験です。このようにすれば、jupyter notebookからMetalにアクセスでき、Convolutionのような時間のかかる処理だけをCPU上の計算からMetal上の計算に切り替えて学習スピードを多少は上げることができるかもしれないという目論見です。
実験内容と結果
まずは簡易的な実験として、複雑な数値計算（項をかなり多くしたフーリエ級数を10万サンプルに対して展開し、それをさらに数値微分する）を、Numpyを使った計算とPython+Metalによる計算の双方に対して同様の条件下で行った際のパフォーマンスを計測することを試みました。実験内容の詳細と結果は下のリンクにまとめています。
結論から言えば、Python+Metalのほうが高速で
■10イテレーションの平均実行時間 (Second)
・Python + Numpy
・平均値: 7.1314005851745605
・中央値: 7.05201256275177
・Python + Metal
・平均値: 0.10262222290039062 （初回起動時のオーバーヘッド考慮）
・中央値: 0.008315801620483398
のようになりました。このようにMetalを利用すると本当に計算できてんのか？っていうくらい爆速でした。求めた結果のcos類似度をnumpyとMetalの計算結果の双方で比べることをしましたが、誤差レベルの相違しか見られずどうやら計算できているみたいです。（これは推測の域を出ないですが、速さの一因として、.metallibにコンパイルする際に、フーリエ級数のためのforループがShaderコンパイラの最適化により事前にインライン化されているのかもしれません。なお、Numpy側のコードはさらなる最適化の余地を残しているものの、単純な演算の繰り返しにおけるGPUの実行速度とそのお手軽さを見せたくてこのようなベンチマークとしています。）
その他に、ベンチマーク用の関数の複雑性を増すことで、Metalでは処理性能の劣化がCPU処理に比べ少ないことも確認できました。
なお、今回はディープラーニングを実行するまではできてないですが、次はCNNにおいてどのくらい学習速度が向上するかの実験を行いたいと思いますので、記事の投稿までしばらくお待ち下さい。ここまでで、今回の簡易実験に関する話はおわりです。以降はPythonからSwift -> Metalを実行する際にどのようにオーバーヘッド無く行うかという技術的なフォローに関することを書いていきますので、興味のない方はここでこの記事を閉じてください（笑）
番外編: PythonからMetalを呼び出すためのガイド
これは少しハッキーで、Swiftをある程度書ける必要があるのですが、まずはプログラムの作成及び、実行フローを簡単に紹介します。SwiftからMetal APIを介してGPU処理を実行するためには
Step.1: Swift Package Managerを利用してdynamic library用のSwift Projectを作成する
Step.2: Metal Shaderのkernel関数を書く
Step.3: SwiftからMetal Shaderを実行する
step1が結構肝でして、この実験ではいわゆるPython -> Swiftをプロセス間通信を通して行うというものではなく、Swiftのプログラムをdylibとしてエクスポートし、それをPython上からctypesによってネイティブに実行するという手法を取っています。（これがハッキーな所以です）実際、プロセス間通信で大量のRGB値をPython -> Swiftに送りそれをSwift上でコピーするみたいなことをしたら大変計算効率が悪いですよね。よって、入力データはコピーしたくないのです。
Swift Package Managerを利用してdynamicライブラリを作成するには、CLIから
$ mkdir SampleProject
$ cd SampleProject
$ swift package init --type libraryを入力し、Package.swiftのlibraryタイプの指定を
のように指定します。
Step.2のMetal ShderはいわゆるC++ライクなDSLであるMetal-Shading-Language を利用して記述します。基本的にはmetal_stdlibにあるAPIを用いてプログラミングするため、従来のCPU向けのiostreamなどはincludeできません。
この言語で書いた Shaders.metal のようなファイルをまずは作成します。
これを building_a_library_with_metal_s_command-line_tools に沿ってビルドします。ビルド後、.metallibファイルが作成されます。
Step.3において、Swift側からこの.metallibをロードし、Metalのプログラムを実行するためのパイプラインを作成します。そのファイルが PyMetalBridge.swift です。
このとき、swift_sigmoid_on_gpuやswift_differential_on_gpuの引数がUnsafePointerとなっておりますが、これはpythonの数値型とSwiftの数値型では内部的な扱いが異なるため、この2つの言語が共通で扱えるCのFloatのポインタとして値をやり取りしています。（ここで注意が必要なのが、Metal(GPU)は倍精度のfloat64をサポートしておらず、32-bit floatで計算されます。よって、Swiftに渡す前に、array.astype("float32")でfloat64をfloat32にキャストする必要があります。）
Python -> Swiftの関数呼び出し
次にPythonからSwiftで作成したMetal実行関数の呼び出しに関して、Pythonにおいては通常のデータ分析やディープラーニングのときのようにnumpyでndarrayを作成し、これをctypesにコンバートしてSwiftに渡してあげればよいだけです。Swiftで実装したsigmoid関数をPythonから呼び出すサンプルが次の画像のとおりです。
ctypes.CDLLでdylibを読み込み、関数pointerに引数を設定します。そして、numpyで作成したarrayを
input_ptr = input_array.ctypes.data_as(ctypes.POINTER(ctypes.c_float))でpointer型に変換します。また、outputの受け取りは
output_mutable_ptr = (ctypes.c_float * len(input_array))()のようにpointerを初期化し、そのpointerを関数の引数として渡します。このoutput_mutable_ptrにMetalでの実行結果が書き込まれます。また、Swift側ではReadOnlyなpointerをUnsafePointer<T>、ReadAndWriteなpointerをUnsafeMutablePointer<T>として扱うので、UnsafePointer<T>を入力に、UnsafeMutablePointer<T>を出力用の型として定義しています。
以上が簡単なPython -> Swift + Metal の実行フローとなります。
どうでしたでしょうか。まだまだ本当に実験的なもので、ディープラーニングに使えそうという結論は出ていないものの、複雑な数値計算においてはMetalでかなりのパフォーマンスが出るという実験結果が出ました。なお、Metalを使ったのは今回が初めてなので、間違っていることなどがあればご指摘いただければ嬉しいです。
スペシャルサンクス
実験にあたって、仲のいいソフトウェアエンジニアである @dealforest さんにMetalのXcode上でのデバッグを手伝ってもらいました。ありがとうございました！
デバッグ時のブログ: http://dealforest.hatenablog.com/entry/2020/06/02/195305