dc-ae-f64c128-in-1.0 - 深度压缩自编码器提高高分辨率模型性能

Deep Compression Autoencoder (DC-AE)项目介绍

项目背景

Deep Compression Autoencoder (DC-AE) 是一项技术前沿的项目，专注于提升高分辨率扩散模型的效率。现有的自编码器通常在中等空间压缩比（例如8倍）的情况下表现出色，但面对高空间压缩比（例如64倍）时，重建精度会大幅下降。DC-AE旨在通过创新的技术提高自编码器的空间压缩比至128倍，同时保持重建质量。

项目亮点

残差自编码技术：DC-AE 引入残差自编码技术，通过设计模型学习空间到通道转换特征的残差，以减少高空间压缩自编码器的优化难度。
解耦高分辨率适应：通过三阶段的训练策略，缓解高空间压缩自编码器的一般化损失。该策略高效地实现了模型的训练和适应。
显著的加速效果：通过将DC-AE应用于潜在扩散模型，可以在不牺牲准确性的前提下实现显著的训练和推理速度提升。例如，在ImageNet 512x512的实验中，DC-AE在H100 GPU上实现了19.1倍的推理速度提升和17.9倍的训练速度提升，同时取得更优的FID分数。

核心技术及应用

深度压缩自编码器（DC-AE）

DC-AE的模型可以通过轻松的代码完成构建和使用。以下是一个简单的使用示例：

from efficientvit.ae_model_zoo import DCAE_HF

dc_ae = DCAE_HF.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0")

# 图像编码和解码
from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.utils import save_image

device = torch.device("cuda")
dc_ae = dc_ae.to(device).eval()

# 图像转换
transform = transforms.Compose([
    transforms.Resize(512),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
image = Image.open("assets/fig/girl.png")
x = transform(image)[None].to(device)
latent = dc_ae.encode(x)
print(latent.shape)

# 图像解码
y = dc_ae.decode(latent)
save_image(y * 0.5 + 0.5, "demo_dc_ae.png")

高效扩散模型与DC-AE的集成

DC-AE不仅仅应用于自编码器，还可以与扩散模型进行良好的集成，从而使得模型在潜在空间进行去噪处理，并保持高效的生成效能：

from efficientvit.diffusion_model_zoo import DCAE_Diffusion_HF

dc_ae_diffusion = DCAE_Diffusion_HF.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0-uvit-h-in-512px-train2000k")

# 在潜在空间进行去噪
import torch
import numpy as np
from torchvision.utils import save_image

torch.set_grad_enabled(False)
device = torch.device("cuda")
dc_ae_diffusion = dc_ae_diffusion.to(device).eval()

latent_samples = dc_ae_diffusion.diffusion_model.generate(prompts, prompts_null, 6.0, eval_generator)
latent_samples = latent_samples / dc_ae_diffusion.scaling_factor

# 图像生成
image_samples = dc_ae_diffusion.autoencoder.decode(latent_samples)
save_image(image_samples * 0.5 + 0.5, "demo_dc_ae_diffusion.png", nrow=int(np.sqrt(num_samples)))

结论

DC-AE通过引入残差自编码和解耦高分辨率适应等创新技术，大幅提升了高空间压缩率下自编码器的性能，使其在高分辨率图像生成任务中发挥重要作用。它不仅提高了图像生成的效率，同时在保持精度方面也取得了显著的突破。对于科学研究和实际应用来说，DC-AE都有极高的价值。