【低显存】 diffusion 爆显存的解决方法

1. 降低分辨率
2. 降低通道数
3. 使用latent diffusion，例如stable
4. 使用低精度训练

下面将分别讲解如何使用这些方法。

使用latent diffusion

目前的主流模型都用了encoder降低中间层resolution，但是有一些比较老的工作，比如sr3，还是老架构。GitHub上有很多改进的工作，可以找找看。不想换的话，只需要配置好diffusers库，然后：
使用图像的低维表示而不是实际的像素空间，这使得它的内存效率更高。编码器将图像压缩为更小的表示，解码器将压缩的表示转换回图像。对于文本到图像模型，您需要一个分词器和一个编码器来生成文本嵌入。SD具有三个独立的预训练模型。

使用**from_pretrained()方法加载所有这些组件。您可以在预训练[runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)**检查点中找到它们，每个组件都存储在单独的子文件夹中

from PIL import Image
import torch
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-5", subfolder="vae", use_safetensors=True)
tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-5", subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(
    "CompVis/stable-diffusion-v1-5", subfolder="text_encoder", use_safetensors=True
)
unet = UNet2DConditionModel.from_pretrained(
    "CompVis/stable-diffusion-v1-5", subfolder="unet", use_safetensors=True
)

代替默认的**PNDMScheduler，将其替换为UniPCMultistepScheduler**以插入不同的调度程序：

from diffusers import UniPCMultistepScheduler
scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

使用低精度推理

参考，照着来即可

https://datawhalechina.github.io/thorough-pytorch/%E7%AC%AC%E5%85%AD%E7%AB%A0/6.4%20%E5%8D%8A%E7%B2%BE%E5%BA%A6%E8%AE%AD%E7%BB%83.html

在PyTorch中使用autocast配置半精度训练，同时需要在下面三处加以设置：

import autocast

from torch.cuda.amp import autocast
模型设置

在模型定义中，使用python的装饰器方法，用autocast装饰模型中的forward函数。关于装饰器的使用，可以参考这里：

@autocast()
def forward(self, x):
…
return x
训练过程

在训练过程中，只需在将数据输入模型及其之后的部分放入“with autocast():“即可：

for x in train_loader:
x = x.cuda()
with autocast():
output = model(x)
…
注意：

半精度训练主要适用于数据本身的size比较大（比如说3D图像、视频等）。当数据本身的size并不大时（比如手写数字MNIST数据集的图片尺寸只有28*28），使用半精度训练则可能不会带来显著的提升。