HuggingFace Diffusers 0.3 : 最適化 : メモリと速度

by

HuggingFace Diffusers 0.3 : 最適化 : メモリと速度 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 09/29/2022 (v0.3.0)

* 本ページは、HuggingFace Diffusers の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

クラスキャット 人工知能 研究開発支援サービス

◆ クラスキャット は人工知能・テレワークに関する各種サービスを提供しています。お気軽にご相談ください :

◆ 人工知能とビジネスをテーマに WEB セミナーを定期的に開催しています。スケジュール。
  • お住まいの地域に関係なく Web ブラウザからご参加頂けます。事前登録 が必要ですのでご注意ください。

◆ お問合せ : 本件に関するお問い合わせ先は下記までお願いいたします。

  • 株式会社クラスキャット セールス・マーケティング本部 セールス・インフォメーション
  • sales-info@classcat.com  ;  Web: www.classcat.com  ;   ClassCatJP

 

 

HuggingFace Diffusers 0.3 : 最適化 : メモリと速度

メモリやスピードに対して 🤗 Diffusers 推論を最適化するための幾つかのテクニックとアイデアを提示します。

 

CUDA autocast

CUDA GPU を使用する場合、(僅かに精度は下がりますが) およそ 2 倍の速さで推論を実行するために torch.autocast を活用できます。必要なことは推論呼び出しを autocast コンテキストマネージャ内に配置することだけです。次のサンプルは、Stable Diffusion テキスト-to-画像生成を例として使用しそれを行なう方法を示します :

from torch import autocast
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=True)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt).images[0]

精度損失にもかかわらず、私たちの経験では、最終的な画像結果は float32 版と同じように見えます。Feel free to experiment and report back!

 

半精度重み

GPU メモリを節約するため、モデル重みを直接半精度でロードすることができます。これは fp16 という名前のブランチにセーブされた、float16 版の重みをロードして、それらをロードするときに PyTorch に float16 型を使用するように伝えます :

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=True
)

 

更なるメモリ節約のための sliced アテンション

更ならメモリ節約のために、アテンションの sliced 版を利用できます、これは計算のすべてを一度に実行する代わりにステップ毎に行います。

アテンション・スライシングは単に 1 のバッチサイズが使用された場合でさえ、モデルが一つより多いアテンションヘッドを使用している限りは有用です。1 つより多いアテンションヘッドがある場合、*QK^T* アテンション行列は各ヘッドに対してシーケンシャルに計算できて、これはメモリのかなりの量を節約できます。

各ヘッドに対してアテンション計算をシーケンシャルに実行するには、推論の前にパイプラインで enable_attention_slicing() を呼び出す必要があるだけです、このようにです :

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16",
    torch_dtype=torch.float16,
    use_auth_token=True
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_attention_slicing()
with torch.autocast("cuda"):
    image = pipe(prompt).images[0]

約 10% 遅い推論時間という小さなパフォーマンス・ペナルティはありますが、この方法は Stable Diffusion を 3.2 GB ほどに少ない VRAM で利用することを可能にします！

 

以上