|
<!--Copyright 2024 The HuggingFace Team. All rights reserved. |
|
|
|
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with |
|
the License. You may obtain a copy of the License at |
|
|
|
http://www.apache.org/licenses/LICENSE-2.0 |
|
|
|
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on |
|
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the |
|
specific language governing permissions and limitations under the License. |
|
--> |
|
|
|
# ๋ฉ๋ชจ๋ฆฌ์ ์๋ |
|
|
|
๋ฉ๋ชจ๋ฆฌ ๋๋ ์๋์ ๋ํด ๐ค Diffusers *์ถ๋ก *์ ์ต์ ํํ๊ธฐ ์ํ ๋ช ๊ฐ์ง ๊ธฐ์ ๊ณผ ์์ด๋์ด๋ฅผ ์ ์ํฉ๋๋ค. |
|
์ผ๋ฐ์ ์ผ๋ก, memory-efficient attention์ ์ํด [xFormers](https://github.com/facebookresearch/xformers) ์ฌ์ฉ์ ์ถ์ฒํ๊ธฐ ๋๋ฌธ์, ์ถ์ฒํ๋ [์ค์น ๋ฐฉ๋ฒ](xformers)์ ๋ณด๊ณ ์ค์นํด ๋ณด์ธ์. |
|
|
|
๋ค์ ์ค์ ์ด ์ฑ๋ฅ๊ณผ ๋ฉ๋ชจ๋ฆฌ์ ๋ฏธ์น๋ ์ํฅ์ ๋ํด ์ค๋ช
ํฉ๋๋ค. |
|
|
|
| | ์ง์ฐ์๊ฐ | ์๋ ํฅ์ | |
|
| ---------------- | ------- | ------- | |
|
| ๋ณ๋ ์ค์ ์์ | 9.50s | x1 | |
|
| cuDNN auto-tuner | 9.37s | x1.01 | |
|
| fp16 | 3.61s | x2.63 | |
|
| Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ | 3.30s | x2.88 | |
|
| traced UNet | 3.21s | x2.96 | |
|
| memory-efficient attention | 2.63s | x3.61 | |
|
|
|
<em> |
|
NVIDIA TITAN RTX์์ 50 DDIM ์คํ
์ "a photo of an astronaut riding a horse on mars" ํ๋กฌํํธ๋ก 512x512 ํฌ๊ธฐ์ ๋จ์ผ ์ด๋ฏธ์ง๋ฅผ ์์ฑํ์์ต๋๋ค. |
|
</em> |
|
|
|
## cuDNN auto-tuner ํ์ฑํํ๊ธฐ |
|
|
|
[NVIDIA cuDNN](https://developer.nvidia.com/cudnn)์ ์ปจ๋ณผ๋ฃจ์
์ ๊ณ์ฐํ๋ ๋ง์ ์๊ณ ๋ฆฌ์ฆ์ ์ง์ํฉ๋๋ค. Autotuner๋ ์งง์ ๋ฒค์น๋งํฌ๋ฅผ ์คํํ๊ณ ์ฃผ์ด์ง ์
๋ ฅ ํฌ๊ธฐ์ ๋ํด ์ฃผ์ด์ง ํ๋์จ์ด์์ ์ต๊ณ ์ ์ฑ๋ฅ์ ๊ฐ์ง ์ปค๋์ ์ ํํฉ๋๋ค. |
|
|
|
**์ปจ๋ณผ๋ฃจ์
๋คํธ์ํฌ**๋ฅผ ํ์ฉํ๊ณ ์๊ธฐ ๋๋ฌธ์ (๋ค๋ฅธ ์ ํ๋ค์ ํ์ฌ ์ง์๋์ง ์์), ๋ค์ ์ค์ ์ ํตํด ์ถ๋ก ์ ์ cuDNN autotuner๋ฅผ ํ์ฑํํ ์ ์์ต๋๋ค: |
|
|
|
```python |
|
import torch |
|
|
|
torch.backends.cudnn.benchmark = True |
|
``` |
|
|
|
### fp32 ๋์ tf32 ์ฌ์ฉํ๊ธฐ (Ampere ๋ฐ ์ดํ CUDA ์ฅ์น๋ค์์) |
|
|
|
Ampere ๋ฐ ์ดํ CUDA ์ฅ์น์์ ํ๋ ฌ๊ณฑ ๋ฐ ์ปจ๋ณผ๋ฃจ์
์ TensorFloat32(TF32) ๋ชจ๋๋ฅผ ์ฌ์ฉํ์ฌ ๋ ๋น ๋ฅด์ง๋ง ์ฝ๊ฐ ๋ ์ ํํ ์ ์์ต๋๋ค. |
|
๊ธฐ๋ณธ์ ์ผ๋ก PyTorch๋ ์ปจ๋ณผ๋ฃจ์
์ ๋ํด TF32 ๋ชจ๋๋ฅผ ํ์ฑํํ์ง๋ง ํ๋ ฌ ๊ณฑ์
์ ํ์ฑํํ์ง ์์ต๋๋ค. |
|
๋คํธ์ํฌ์ ์์ ํ float32 ์ ๋ฐ๋๊ฐ ํ์ํ ๊ฒฝ์ฐ๊ฐ ์๋๋ฉด ํ๋ ฌ ๊ณฑ์
์ ๋ํด์๋ ์ด ์ค์ ์ ํ์ฑํํ๋ ๊ฒ์ด ์ข์ต๋๋ค. |
|
์ด๋ ์ผ๋ฐ์ ์ผ๋ก ๋ฌด์ํ ์ ์๋ ์์น์ ์ ํ๋ ์์ค์ด ์์ง๋ง, ๊ณ์ฐ ์๋๋ฅผ ํฌ๊ฒ ๋์ผ ์ ์์ต๋๋ค. |
|
๊ทธ๊ฒ์ ๋ํด [์ฌ๊ธฐ](https://huggingface.co/docs/transformers/v4.18.0/en/performance#tf32)์ ๋ ์ฝ์ ์ ์์ต๋๋ค. |
|
์ถ๋ก ํ๊ธฐ ์ ์ ๋ค์์ ์ถ๊ฐํ๊ธฐ๋ง ํ๋ฉด ๋ฉ๋๋ค: |
|
|
|
```python |
|
import torch |
|
|
|
torch.backends.cuda.matmul.allow_tf32 = True |
|
``` |
|
|
|
## ๋ฐ์ ๋ฐ๋ ๊ฐ์ค์น |
|
|
|
๋ ๋ง์ GPU ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ๊ณ ๋ ๋น ๋ฅธ ์๋๋ฅผ ์ป๊ธฐ ์ํด ๋ชจ๋ธ ๊ฐ์ค์น๋ฅผ ๋ฐ์ ๋ฐ๋(half precision)๋ก ์ง์ ๋ถ๋ฌ์ค๊ณ ์คํํ ์ ์์ต๋๋ค. |
|
์ฌ๊ธฐ์๋ `fp16`์ด๋ผ๋ ๋ธ๋์น์ ์ ์ฅ๋ float16 ๋ฒ์ ์ ๊ฐ์ค์น๋ฅผ ๋ถ๋ฌ์ค๊ณ , ๊ทธ ๋ `float16` ์ ํ์ ์ฌ์ฉํ๋๋ก PyTorch์ ์ง์ํ๋ ์์
์ด ํฌํจ๋ฉ๋๋ค. |
|
|
|
```Python |
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/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] |
|
``` |
|
|
|
<Tip warning={true}> |
|
์ด๋ค ํ์ดํ๋ผ์ธ์์๋ [`torch.autocast`](https://pytorch.org/docs/stable/amp.html#torch.autocast) ๋ฅผ ์ฌ์ฉํ๋ ๊ฒ์ ๊ฒ์์ ์ด๋ฏธ์ง๋ฅผ ์์ฑํ ์ ์๊ณ , ์์ํ float16 ์ ๋ฐ๋๋ฅผ ์ฌ์ฉํ๋ ๊ฒ๋ณด๋ค ํญ์ ๋๋ฆฌ๊ธฐ ๋๋ฌธ์ ์ฌ์ฉํ์ง ์๋ ๊ฒ์ด ์ข์ต๋๋ค. |
|
</Tip> |
|
|
|
## ์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ ์ฌ๋ผ์ด์ค ์ดํ
์
|
|
|
|
์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด, ํ ๋ฒ์ ๋ชจ๋ ๊ณ์ฐํ๋ ๋์ ๋จ๊ณ์ ์ผ๋ก ๊ณ์ฐ์ ์ํํ๋ ์ฌ๋ผ์ด์ค ๋ฒ์ ์ ์ดํ
์
(attention)์ ์ฌ์ฉํ ์ ์์ต๋๋ค. |
|
|
|
<Tip> |
|
Attention slicing์ ๋ชจ๋ธ์ด ํ๋ ์ด์์ ์ดํ
์
ํค๋๋ฅผ ์ฌ์ฉํ๋ ํ, ๋ฐฐ์น ํฌ๊ธฐ๊ฐ 1์ธ ๊ฒฝ์ฐ์๋ ์ ์ฉํฉ๋๋ค. |
|
ํ๋ ์ด์์ ์ดํ
์
ํค๋๊ฐ ์๋ ๊ฒฝ์ฐ *QK^T* ์ดํ
์
๋งคํธ๋ฆญ์ค๋ ์๋นํ ์์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ ์ ์๋ ๊ฐ ํค๋์ ๋ํด ์์ฐจ์ ์ผ๋ก ๊ณ์ฐ๋ ์ ์์ต๋๋ค. |
|
</Tip> |
|
|
|
๊ฐ ํค๋์ ๋ํด ์์ฐจ์ ์ผ๋ก ์ดํ
์
๊ณ์ฐ์ ์ํํ๋ ค๋ฉด, ๋ค์๊ณผ ๊ฐ์ด ์ถ๋ก ์ ์ ํ์ดํ๋ผ์ธ์์ [`~StableDiffusionPipeline.enable_attention_slicing`]๋ฅผ ํธ์ถํ๋ฉด ๋ฉ๋๋ค: |
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
|
|
torch_dtype=torch.float16, |
|
) |
|
pipe = pipe.to("cuda") |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_attention_slicing() |
|
image = pipe(prompt).images[0] |
|
``` |
|
|
|
์ถ๋ก ์๊ฐ์ด ์ฝ 10% ๋๋ ค์ง๋ ์ฝ๊ฐ์ ์ฑ๋ฅ ์ ํ๊ฐ ์์ง๋ง ์ด ๋ฐฉ๋ฒ์ ์ฌ์ฉํ๋ฉด 3.2GB ์ ๋์ ์์ VRAM์ผ๋ก๋ Stable Diffusion์ ์ฌ์ฉํ ์ ์์ต๋๋ค! |
|
|
|
|
|
## ๋ ํฐ ๋ฐฐ์น๋ฅผ ์ํ sliced VAE ๋์ฝ๋ |
|
|
|
์ ํ๋ VRAM์์ ๋๊ท๋ชจ ์ด๋ฏธ์ง ๋ฐฐ์น๋ฅผ ๋์ฝ๋ฉํ๊ฑฐ๋ 32๊ฐ ์ด์์ ์ด๋ฏธ์ง๊ฐ ํฌํจ๋ ๋ฐฐ์น๋ฅผ ํ์ฑํํ๊ธฐ ์ํด, ๋ฐฐ์น์ latent ์ด๋ฏธ์ง๋ฅผ ํ ๋ฒ์ ํ๋์ฉ ๋์ฝ๋ฉํ๋ ์ฌ๋ผ์ด์ค VAE ๋์ฝ๋๋ฅผ ์ฌ์ฉํ ์ ์์ต๋๋ค. |
|
|
|
์ด๋ฅผ [`~StableDiffusionPipeline.enable_attention_slicing`] ๋๋ [`~StableDiffusionPipeline.enable_xformers_memory_efficient_attention`]๊ณผ ๊ฒฐํฉํ์ฌ ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ์ ์ถ๊ฐ๋ก ์ต์ํํ ์ ์์ต๋๋ค. |
|
|
|
VAE ๋์ฝ๋๋ฅผ ํ ๋ฒ์ ํ๋์ฉ ์ํํ๋ ค๋ฉด ์ถ๋ก ์ ์ ํ์ดํ๋ผ์ธ์์ [`~StableDiffusionPipeline.enable_vae_slicing`]์ ํธ์ถํฉ๋๋ค. ์๋ฅผ ๋ค์ด: |
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
|
|
torch_dtype=torch.float16, |
|
) |
|
pipe = pipe.to("cuda") |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_vae_slicing() |
|
images = pipe([prompt] * 32).images |
|
``` |
|
|
|
๋ค์ค ์ด๋ฏธ์ง ๋ฐฐ์น์์ VAE ๋์ฝ๋๊ฐ ์ฝ๊ฐ์ ์ฑ๋ฅ ํฅ์์ด ์ด๋ฃจ์ด์ง๋๋ค. ๋จ์ผ ์ด๋ฏธ์ง ๋ฐฐ์น์์๋ ์ฑ๋ฅ ์ํฅ์ ์์ต๋๋ค. |
|
|
|
|
|
<a name="sequential_offloading"></a> |
|
## ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด ๊ฐ์ ๊ธฐ๋ฅ์ ์ฌ์ฉํ์ฌ CPU๋ก ์คํ๋ก๋ฉ |
|
|
|
์ถ๊ฐ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํด ๊ฐ์ค์น๋ฅผ CPU๋ก ์คํ๋ก๋ํ๊ณ ์๋ฐฉํฅ ์ ๋ฌ์ ์ํํ ๋๋ง GPU๋ก ๋ก๋ํ ์ ์์ต๋๋ค. |
|
|
|
CPU ์คํ๋ก๋ฉ์ ์ํํ๋ ค๋ฉด [`~StableDiffusionPipeline.enable_sequential_cpu_offload`]๋ฅผ ํธ์ถํ๊ธฐ๋ง ํ๋ฉด ๋ฉ๋๋ค: |
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
|
|
torch_dtype=torch.float16, |
|
) |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_sequential_cpu_offload() |
|
image = pipe(prompt).images[0] |
|
``` |
|
|
|
๊ทธ๋ฌ๋ฉด ๋ฉ๋ชจ๋ฆฌ ์๋น๋ฅผ 3GB ๋ฏธ๋ง์ผ๋ก ์ค์ผ ์ ์์ต๋๋ค. |
|
|
|
์ฐธ๊ณ ๋ก ์ด ๋ฐฉ๋ฒ์ ์ ์ฒด ๋ชจ๋ธ์ด ์๋ ์๋ธ๋ชจ๋ ์์ค์์ ์๋ํฉ๋๋ค. ์ด๋ ๋ฉ๋ชจ๋ฆฌ ์๋น๋ฅผ ์ต์ํํ๋ ๊ฐ์ฅ ์ข์ ๋ฐฉ๋ฒ์ด์ง๋ง ํ๋ก์ธ์ค์ ๋ฐ๋ณต์ ํน์ฑ์ผ๋ก ์ธํด ์ถ๋ก ์๋๊ฐ ํจ์ฌ ๋๋ฆฝ๋๋ค. ํ์ดํ๋ผ์ธ์ UNet ๊ตฌ์ฑ ์์๋ ์ฌ๋ฌ ๋ฒ ์คํ๋ฉ๋๋ค('num_inference_steps' ๋งํผ). ๋งค๋ฒ UNet์ ์๋ก ๋ค๋ฅธ ์๋ธ๋ชจ๋์ด ์์ฐจ์ ์ผ๋ก ์จ๋ก๋๋ ๋ค์ ํ์์ ๋ฐ๋ผ ์คํ๋ก๋๋๋ฏ๋ก ๋ฉ๋ชจ๋ฆฌ ์ด๋ ํ์๊ฐ ๋ง์ต๋๋ค. |
|
|
|
<Tip> |
|
๋ ๋ค๋ฅธ ์ต์ ํ ๋ฐฉ๋ฒ์ธ <a href="#model_offloading">๋ชจ๋ธ ์คํ๋ก๋ฉ</a>์ ์ฌ์ฉํ๋ ๊ฒ์ ๊ณ ๋ คํ์ญ์์ค. ์ด๋ ํจ์ฌ ๋น ๋ฅด์ง๋ง ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ด ํฌ์ง๋ ์์ต๋๋ค. |
|
</Tip> |
|
|
|
๋ํ ttention slicing๊ณผ ์ฐ๊ฒฐํด์ ์ต์ ๋ฉ๋ชจ๋ฆฌ(< 2GB)๋ก๋ ๋์ํ ์ ์์ต๋๋ค. |
|
|
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
|
|
torch_dtype=torch.float16, |
|
) |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_sequential_cpu_offload() |
|
pipe.enable_attention_slicing(1) |
|
|
|
image = pipe(prompt).images[0] |
|
``` |
|
|
|
**์ฐธ๊ณ **: 'enable_sequential_cpu_offload()'๋ฅผ ์ฌ์ฉํ ๋, ๋ฏธ๋ฆฌ ํ์ดํ๋ผ์ธ์ CUDA๋ก ์ด๋ํ์ง **์๋** ๊ฒ์ด ์ค์ํฉ๋๋ค.๊ทธ๋ ์ง ์์ผ๋ฉด ๋ฉ๋ชจ๋ฆฌ ์๋น์ ์ด๋์ด ์ต์ํ๋ฉ๋๋ค. ๋ ๋ง์ ์ ๋ณด๋ฅผ ์ํด [์ด ์ด์](https://github.com/huggingface/diffusers/issues/1934)๋ฅผ ๋ณด์ธ์. |
|
|
|
<a name="model_offloading"></a> |
|
## ๋น ๋ฅธ ์ถ๋ก ๊ณผ ๋ฉ๋ชจ๋ฆฌ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ ๋ชจ๋ธ ์คํ๋ก๋ฉ |
|
|
|
[์์ฐจ์ CPU ์คํ๋ก๋ฉ](#sequential_offloading)์ ์ด์ ์น์
์์ ์ค๋ช
ํ ๊ฒ์ฒ๋ผ ๋ง์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ๋ณด์กดํ์ง๋ง ํ์์ ๋ฐ๋ผ ์๋ธ๋ชจ๋์ GPU๋ก ์ด๋ํ๊ณ ์ ๋ชจ๋์ด ์คํ๋ ๋ ์ฆ์ CPU๋ก ๋ฐํ๋๊ธฐ ๋๋ฌธ์ ์ถ๋ก ์๋๊ฐ ๋๋ ค์ง๋๋ค. |
|
|
|
์ ์ฒด ๋ชจ๋ธ ์คํ๋ก๋ฉ์ ๊ฐ ๋ชจ๋ธ์ ๊ตฌ์ฑ ์์์ธ _modules_์ ์ฒ๋ฆฌํ๋ ๋์ , ์ ์ฒด ๋ชจ๋ธ์ GPU๋ก ์ด๋ํ๋ ๋์์
๋๋ค. ์ด๋ก ์ธํด ์ถ๋ก ์๊ฐ์ ๋ฏธ์น๋ ์ํฅ์ ๋ฏธ๋ฏธํ์ง๋ง(ํ์ดํ๋ผ์ธ์ 'cuda'๋ก ์ด๋ํ๋ ๊ฒ๊ณผ ๋น๊ตํ์ฌ) ์ฌ์ ํ ์ฝ๊ฐ์ ๋ฉ๋ชจ๋ฆฌ๋ฅผ ์ ์ฝํ ์ ์์ต๋๋ค. |
|
|
|
์ด ์๋๋ฆฌ์ค์์๋ ํ์ดํ๋ผ์ธ์ ์ฃผ์ ๊ตฌ์ฑ ์์ ์ค ํ๋๋ง(์ผ๋ฐ์ ์ผ๋ก ํ
์คํธ ์ธ์ฝ๋, unet ๋ฐ vae) GPU์ ์๊ณ , ๋๋จธ์ง๋ CPU์์ ๋๊ธฐํ ๊ฒ์
๋๋ค. |
|
์ฌ๋ฌ ๋ฐ๋ณต์ ์ํด ์คํ๋๋ UNet๊ณผ ๊ฐ์ ๊ตฌ์ฑ ์์๋ ๋ ์ด์ ํ์ํ์ง ์์ ๋๊น์ง GPU์ ๋จ์ ์์ต๋๋ค. |
|
|
|
์ด ๊ธฐ๋ฅ์ ์๋์ ๊ฐ์ด ํ์ดํ๋ผ์ธ์์ `enable_model_cpu_offload()`๋ฅผ ํธ์ถํ์ฌ ํ์ฑํํ ์ ์์ต๋๋ค. |
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
torch_dtype=torch.float16, |
|
) |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_model_cpu_offload() |
|
image = pipe(prompt).images[0] |
|
``` |
|
|
|
์ด๋ ์ถ๊ฐ์ ์ธ ๋ฉ๋ชจ๋ฆฌ ์ ์ฝ์ ์ํ attention slicing๊ณผ๋ ํธํ๋ฉ๋๋ค. |
|
|
|
```Python |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
torch_dtype=torch.float16, |
|
) |
|
|
|
prompt = "a photo of an astronaut riding a horse on mars" |
|
pipe.enable_model_cpu_offload() |
|
pipe.enable_attention_slicing(1) |
|
|
|
image = pipe(prompt).images[0] |
|
``` |
|
|
|
<Tip> |
|
์ด ๊ธฐ๋ฅ์ ์ฌ์ฉํ๋ ค๋ฉด 'accelerate' ๋ฒ์ 0.17.0 ์ด์์ด ํ์ํฉ๋๋ค. |
|
</Tip> |
|
|
|
## Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ ์ฌ์ฉํ๊ธฐ |
|
|
|
Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์์ ์ฐจ์ ์์๋ฅผ ๋ณด์กดํ๋ ๋ฉ๋ชจ๋ฆฌ์์ NCHW ํ
์ ๋ฐฐ์ด์ ๋์ฒดํ๋ ๋ฐฉ๋ฒ์
๋๋ค. |
|
Channels Last ํ
์๋ ์ฑ๋์ด ๊ฐ์ฅ ์กฐ๋ฐํ ์ฐจ์์ด ๋๋ ๋ฐฉ์์ผ๋ก ์ ๋ ฌ๋ฉ๋๋ค(์ผ๋ช
ํฝ์
๋น ์ด๋ฏธ์ง๋ฅผ ์ ์ฅ). |
|
ํ์ฌ ๋ชจ๋ ์ฐ์ฐ์ Channels Last ํ์์ ์ง์ํ๋ ๊ฒ์ ์๋๋ผ ์ฑ๋ฅ์ด ์ ํ๋ ์ ์์ผ๋ฏ๋ก, ์ฌ์ฉํด๋ณด๊ณ ๋ชจ๋ธ์ ์ ์๋ํ๋์ง ํ์ธํ๋ ๊ฒ์ด ์ข์ต๋๋ค. |
|
|
|
|
|
์๋ฅผ ๋ค์ด ํ์ดํ๋ผ์ธ์ UNet ๋ชจ๋ธ์ด channels Last ํ์์ ์ฌ์ฉํ๋๋ก ์ค์ ํ๋ ค๋ฉด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค: |
|
|
|
```python |
|
print(pipe.unet.conv_out.state_dict()["weight"].stride()) # (2880, 9, 3, 1) |
|
pipe.unet.to(memory_format=torch.channels_last) # in-place ์ฐ์ฐ |
|
# 2๋ฒ์งธ ์ฐจ์์์ ์คํธ๋ผ์ด๋ 1์ ๊ฐ์ง๋ (2880, 1, 960, 320)๋ก, ์ฐ์ฐ์ด ์๋ํจ์ ์ฆ๋ช
ํฉ๋๋ค. |
|
print(pipe.unet.conv_out.state_dict()["weight"].stride()) |
|
``` |
|
|
|
## ์ถ์ (tracing) |
|
|
|
์ถ์ ์ ๋ชจ๋ธ์ ํตํด ์์ ์
๋ ฅ ํ
์๋ฅผ ํตํด ์คํ๋๋๋ฐ, ํด๋น ์
๋ ฅ์ด ๋ชจ๋ธ์ ๋ ์ด์ด๋ฅผ ํต๊ณผํ ๋ ํธ์ถ๋๋ ์์
์ ์บก์ฒํ์ฌ ์คํ ํ์ผ ๋๋ 'ScriptFunction'์ด ๋ฐํ๋๋๋ก ํ๊ณ , ์ด๋ just-in-time ์ปดํ์ผ๋ก ์ต์ ํ๋ฉ๋๋ค. |
|
|
|
UNet ๋ชจ๋ธ์ ์ถ์ ํ๊ธฐ ์ํด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค: |
|
|
|
```python |
|
import time |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
import functools |
|
|
|
# torch ๊ธฐ์ธ๊ธฐ ๋นํ์ฑํ |
|
torch.set_grad_enabled(False) |
|
|
|
# ๋ณ์ ์ค์ |
|
n_experiments = 2 |
|
unet_runs_per_experiment = 50 |
|
|
|
|
|
# ์
๋ ฅ ๋ถ๋ฌ์ค๊ธฐ |
|
def generate_inputs(): |
|
sample = torch.randn((2, 4, 64, 64), device="cuda", dtype=torch.float16) |
|
timestep = torch.rand(1, device="cuda", dtype=torch.float16) * 999 |
|
encoder_hidden_states = torch.randn((2, 77, 768), device="cuda", dtype=torch.float16) |
|
return sample, timestep, encoder_hidden_states |
|
|
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
torch_dtype=torch.float16, |
|
).to("cuda") |
|
unet = pipe.unet |
|
unet.eval() |
|
unet.to(memory_format=torch.channels_last) # Channels Last ๋ฉ๋ชจ๋ฆฌ ํ์ ์ฌ์ฉ |
|
unet.forward = functools.partial(unet.forward, return_dict=False) # return_dict=False์ ๊ธฐ๋ณธ๊ฐ์ผ๋ก ์ค์ |
|
|
|
# ์๋ฐ์
|
|
for _ in range(3): |
|
with torch.inference_mode(): |
|
inputs = generate_inputs() |
|
orig_output = unet(*inputs) |
|
|
|
# ์ถ์ |
|
print("tracing..") |
|
unet_traced = torch.jit.trace(unet, inputs) |
|
unet_traced.eval() |
|
print("done tracing") |
|
|
|
|
|
# ์๋ฐ์
๋ฐ ๊ทธ๋ํ ์ต์ ํ |
|
for _ in range(5): |
|
with torch.inference_mode(): |
|
inputs = generate_inputs() |
|
orig_output = unet_traced(*inputs) |
|
|
|
|
|
# ๋ฒค์น๋งํน |
|
with torch.inference_mode(): |
|
for _ in range(n_experiments): |
|
torch.cuda.synchronize() |
|
start_time = time.time() |
|
for _ in range(unet_runs_per_experiment): |
|
orig_output = unet_traced(*inputs) |
|
torch.cuda.synchronize() |
|
print(f"unet traced inference took {time.time() - start_time:.2f} seconds") |
|
for _ in range(n_experiments): |
|
torch.cuda.synchronize() |
|
start_time = time.time() |
|
for _ in range(unet_runs_per_experiment): |
|
orig_output = unet(*inputs) |
|
torch.cuda.synchronize() |
|
print(f"unet inference took {time.time() - start_time:.2f} seconds") |
|
|
|
# ๋ชจ๋ธ ์ ์ฅ |
|
unet_traced.save("unet_traced.pt") |
|
``` |
|
|
|
๊ทธ ๋ค์, ํ์ดํ๋ผ์ธ์ `unet` ํน์ฑ์ ๋ค์๊ณผ ๊ฐ์ด ์ถ์ ๋ ๋ชจ๋ธ๋ก ๋ฐ๊ฟ ์ ์์ต๋๋ค. |
|
|
|
```python |
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
from dataclasses import dataclass |
|
|
|
|
|
@dataclass |
|
class UNet2DConditionOutput: |
|
sample: torch.Tensor |
|
|
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
torch_dtype=torch.float16, |
|
).to("cuda") |
|
|
|
# jitted unet ์ฌ์ฉ |
|
unet_traced = torch.jit.load("unet_traced.pt") |
|
|
|
|
|
# pipe.unet ์ญ์ |
|
class TracedUNet(torch.nn.Module): |
|
def __init__(self): |
|
super().__init__() |
|
self.in_channels = pipe.unet.config.in_channels |
|
self.device = pipe.unet.device |
|
|
|
def forward(self, latent_model_input, t, encoder_hidden_states): |
|
sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0] |
|
return UNet2DConditionOutput(sample=sample) |
|
|
|
|
|
pipe.unet = TracedUNet() |
|
|
|
with torch.inference_mode(): |
|
image = pipe([prompt] * 1, num_inference_steps=50).images[0] |
|
``` |
|
|
|
|
|
## Memory-efficient attention |
|
|
|
์ดํ
์
๋ธ๋ก์ ๋์ญํญ์ ์ต์ ํํ๋ ์ต๊ทผ ์์
์ผ๋ก GPU ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋์ด ํฌ๊ฒ ํฅ์๋๊ณ ํฅ์๋์์ต๋๋ค. |
|
@tridao์ ๊ฐ์ฅ ์ต๊ทผ์ ํ๋์ ์ดํ
์
: [code](https://github.com/HazyResearch/flash-attention), [paper](https://arxiv.org/pdf/2205.14135.pdf). |
|
|
|
๋ฐฐ์น ํฌ๊ธฐ 1(ํ๋กฌํํธ 1๊ฐ)์ 512x512 ํฌ๊ธฐ๋ก ์ถ๋ก ์ ์คํํ ๋ ๋ช ๊ฐ์ง Nvidia GPU์์ ์ป์ ์๋ ํฅ์์ ๋ค์๊ณผ ๊ฐ์ต๋๋ค: |
|
|
|
| GPU | ๊ธฐ์ค ์ดํ
์
FP16 | ๋ฉ๋ชจ๋ฆฌ ํจ์จ์ ์ธ ์ดํ
์
FP16 | |
|
|------------------ |--------------------- |--------------------------------- | |
|
| NVIDIA Tesla T4 | 3.5it/s | 5.5it/s | |
|
| NVIDIA 3060 RTX | 4.6it/s | 7.8it/s | |
|
| NVIDIA A10G | 8.88it/s | 15.6it/s | |
|
| NVIDIA RTX A6000 | 11.7it/s | 21.09it/s | |
|
| NVIDIA TITAN RTX | 12.51it/s | 18.22it/s | |
|
| A100-SXM4-40GB | 18.6it/s | 29.it/s | |
|
| A100-SXM-80GB | 18.7it/s | 29.5it/s | |
|
|
|
์ด๋ฅผ ํ์ฉํ๋ ค๋ฉด ๋ค์์ ๋ง์กฑํด์ผ ํฉ๋๋ค: |
|
- PyTorch > 1.12 |
|
- Cuda ์ฌ์ฉ ๊ฐ๋ฅ |
|
- [xformers ๋ผ์ด๋ธ๋ฌ๋ฆฌ๋ฅผ ์ค์นํจ](xformers) |
|
```python |
|
from diffusers import StableDiffusionPipeline |
|
import torch |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"stable-diffusion-v1-5/stable-diffusion-v1-5", |
|
torch_dtype=torch.float16, |
|
).to("cuda") |
|
|
|
pipe.enable_xformers_memory_efficient_attention() |
|
|
|
with torch.inference_mode(): |
|
sample = pipe("a small cat") |
|
|
|
# ์ ํ: ์ด๋ฅผ ๋นํ์ฑํ ํ๊ธฐ ์ํด ๋ค์์ ์ฌ์ฉํ ์ ์์ต๋๋ค. |
|
# pipe.disable_xformers_memory_efficient_attention() |
|
``` |
|
|