HuggingFace Diffusers 0.12 : 最適化 : PyTorch 2.0 サポート (翻訳/解説)

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

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

• Accelerated PyTorch 2.0 support in Diffusers

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

HuggingFace Diffusers 0.12 : 最適化 : PyTorch 2.0 サポート

バージョン 0.13.0 から、Diffusers は次期 PyTorch 2.0 リリースからの最新の最適化をサポートしています。これらは以下を含みます :

1. メモリ効率的なアテンションを含む高速化 transformers 実装をサポートします – 追加の依存関係は必要ありません。
2. 個々のモデルがコンパイルされたとき追加の性能ブーストのための torch.compile サポート。

インストール

高速化アテンション実装と torch.compile の恩恵を受けるには、pip から PyTorch 2.0 の最新版をインストールして diffusers 0.13.0 またそれ以上で作業していることを確認するだけです。以下で説明するように、diffusers は利用可能であれば自動的にアテンション最適化を使用します (但し torch.compile はそうではありません)。

pip install --upgrade torch torchvision diffusers

高速化 transformers と torch.compile の使用

1. 高速化 Transformers の実装

PyTorch 2.0 は torch.nn.functional.scaled_dot_product_attention 関数を通して最適化されたメモリ効率的なアテンション実装を含みます、これは入力と GPU タイプに応じて幾つかの最適化を自動的に有効にします。これは xFormers の memory_efficient_attention に類似していますが、PyTorch にネイティブで組み込まれています。

これらの最適化は、PyTorch 2.0 がインストールされていて torch.nn.functional.scaled_dot_product_attention が利用可能であれば、Diffusers でデフォルトで有効にされます。それを使用するには、上記で提示されているように torch 2.0 をインストールするして単純にパイプラインを使用するだけです。例えば :

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

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

If you want to enable it explicitly (which is not required), you can do so as shown below.

import torch
from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import AttnProcessor2_0

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipe.unet.set_attn_processor(AttnProcessor2_0())

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

This should be as fast and memory efficient as xFormers. More details in our benchmark.

torch.compile

追加のスピードアップを得るには、新しい torch.compile 機能を使用することができます。そのためには、unet を torch.compile で単純にラップします。詳細と様々なオプションについては、torch compile docs を参照してください。

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipe.unet = torch.compile(pipe.unet)

batch_size = 10
prompt = "A photo of an astronaut riding a horse on marse."
images = pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images

GPU のタイプに応じて、compile() は高速化 transformer 最適化に対して 2-9% の追加のスピードアップをもたらします。ただし、コンパイルは Ampere (A100, 3090), Ada (4090) と Hopper (H100) のようなより最近の GPU アーキテクチャでより多くのパフォーマンス向上を絞り出せることに注意してください。

コンパイルは完了まである程度の時間がかかりますので、パイプラインを一度準備してから同じタイプの推論演算を複数回実行する必要があるような状況に最善に適します。

Benchmark

(訳注: 原文 参照)

以上