画像生成AI「Stable Diffusion」がTensorFlowとKerasCVで約30％高速になることが判明

機械学習プラットフォームのTensorFlowで実行される、Pythonで記述されたディープラーニング用APIが「Keras」です。このKerasを拡張して画像分類、物体検出、画像分割、画像データ補強などを行うためのモジュール式ビルディングブロック「KerasCV」を使うと、画像生成AIのStable Diffusionが約30％高速になるという報告が、KerasCVの開発者による研究チームからあがっています。

High-performance image generation using Stable Diffusion in KerasCV
https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/

Stable Diffusionは2022年8月に一般公開されたオープンソースの画像生成AIで、入力したキーワード(プロンプト)に沿った内容の写真や絵を自動で生成することができます。

キーワードに沿ってまるで人間が描いたような絵や写真を生み出すAI「Stable Diffusion」が一般公開されたので使ってみた - GIGAZINE

Stable Diffusionは、データセットから学習した内容をいったんノイズに戻し、与えられたプロンプトや画像によって方向付けされながらノイズ除去を行って画像を生成する仕組みです。おおまかな仕組みについては、以下の記事を読むとわかります。

画像生成AI「Stable Diffusion」を使いこなすために知っておくと理解が進む「どうやって絵を描いているのか」をわかりやすく図解 - GIGAZINE

もともとStable Diffusionの基本となっているのは、超解像技術です。例えば低解像度の画像を単純に拡大した時、サイズは大きくなっても解像度はそのままなので、「粗さ」が目立ってしまいます。そこで、機械学習アルゴリズムによってこの「粗さ」というノイズを除去して解像度を上げることは以前から行われてきました。

この超解像技術をさらに押し進めて、「超解像で使うモデルに完全なノイズを読み込ませると何を生成するのか？」という疑問からスタートしたのが、Stable Diffusionをはじめとする潜在拡散モデルです。

そこで研究チームは、PyTorch/Hugging Face Diffusersで動作するStable DiffusionをTensorFlow/KerasCVでも実装し、生成ステップ数50で画像を生成した場合の実行時間を計測して比較しました。ウォームスタートでの結果は以下の通り。いずれにしてもKerasCVに実装した方が高速に処理されており、特にTesla T4をGPUに使用した場合は、Hugging Face Diffusersと比較して30％の改善が見られています。

また、コールドスタートでの計測結果は以下の通りで、KerasCVの方が処理に時間がかかってしまう結果になりました。しかし、コールドスタートは最初の1回だけなので、実際に何回もStable Diffusionを動作させる本番環境では無視できると研究チームは述べています。

また、TensorFlowには XLA(Accelerated Linear Algebra)コンパイラが組み込まれています。このXLAコンパイラを使用することで、さらに速度が大幅に向上するとのこと。通常だと画像生成に約8.17秒かかっていたのが、XLAコンパイラのみを使うと約6.26秒に縮まるそうです。

さらに、KerasCVではfloat16の精度で計算を行い、float32形式で重みを保存する「混合精度」を簡単に有効化できるとのこと。この混合精度とXLAコンパイラを組み合わせると、画像生成時間がさらに約4.25秒にまで縮まると研究チームは報告しています。

なお、TensorFlowとKerasCVで動かせるStable Diffusionのリポジトリは、GitHubで公開されています。

GitHub - divamgupta/stable-diffusion-tensorflow: Stable Diffusion in TensorFlow / Keras
https://github.com/divamgupta/stable-diffusion-tensorflow