VQGAN原理与PyTorch实战：从图像离散化到文本生成高清图像

VQGAN大概是这两年生成模型里最被低估的基础组件之一。它的论文出自2021年的CVPR，题目叫《Taming Transformers for High-Resolution Image Synthesis》，讲的是如何用Transformer在图像离散表征上做自回归生成，再配合CNN解码器拿到高清图像。这套思路后来被DALL·E等模型吸收了，如果你把Patch-level的codebook换成现代扩散模型里的隐空间，整条技术路线都是相通的。这篇文章我会用PyTorch作为主框架，带你完整走一遍VQGAN从原理到实现的流程，包括环境搭建、两阶段训练、预训练模型加载，以及如何把文本条件接进Transformer实现文本到图像生成。

我默认你的电脑有一块NVIDIA显卡，哪怕只有8GB显存也能跑通大部分示例。你不需要写过完整的GAN，只要懂基础的CNN和Transformer结构就行。读完之后的收获不会只是“会跑通demo”，而是能理解VQGAN为什么要把图像切成token来建模，以及后续的DALL·E、Stable Diffusion这些模型到底借鉴了它的哪些设计。

1. 为什么要用VQGAN做高清图像生成

传统GAN在处理高分辨率图像时有一个天然矛盾：判别器看到的是下采样后的全局图像，生成器如果想输出256x256甚至512x512的细节，单靠卷积堆叠很难保持全局结构一致。人类看一张猫的图片，不仅要看清毛发纹理，还要求眼睛、鼻子、嘴巴的相对位置不歪，这种全局依赖关系恰恰是CNN不擅长的。

Transformer擅长建模长距离依赖，但把它直接用到像素空间不现实。一张256x256的图像有65536个像素，如果按像素序列做自回归，序列长度太长，计算量会呈平方级爆炸，训练成本高到离谱。VQGAN的处理方式很聪明：先用CNN把图像压缩成一个尺寸小得多的离散token序列，再在这个token序列上跑Transformer。这样序列长度可能从65536降到256，Transformer的全局建模能力正好用在刀刃上。

这其实就是两阶段的思路，第一阶段学一个“图像编码器”，把像素变成有语义的离散token；第二阶段学一个“生成模型”，负责预测token序列。第一阶段对应VQGAN，第二阶段对应Transformer。文本到图像生成时，Transformer接收文本条件，输出完整的图像token序列，最后交给VQGAN的解码器还原成高清图像。整个过程看起来复杂，但每一步都能用PyTorch代码清晰表达。

这套架构的另一大价值是可复用性。第一阶段训练好的VQGAN是通用的，你可以换不同的第二阶段模型，比如换成扩散模型，或者换成别的自回归模型，不需要重新训练图像编码器。所以VQGAN不只是一个具体的模型，更是一种“用离散token连接视觉与语言”的设计范式。

2. VQGAN的核心组件与两阶段训练思想

2.1 编码器、解码器与离散码本

VQGAN的第一阶段本质上是一个带码本约束的自编码器。输入图像经过编码器E得到特征图，量化层再将每个位置的特征向量替换为码本中最接近的向量，解码器G再把这个量化后的特征图还原成图像。

码本是整个模型的灵魂。它是一个可学习的向量表，假设大小为K，每个向量维度是n_z，对应K个“视觉词汇”。量化操作其实就是一个最近邻查找：对特征图上每个位置的特征向量z_e，计算它与码本中所有向量的距离，取距离最近的那个向量z_q作为输出。这里要注意的是，量化本身没有梯度，训练时直接用了一个直通估计器，把解码器传来的梯度原样复制回编码器，让编码器能继续更新。

码本更新也值得细说。最直觉的方式是用梯度下降更新码本向量，但实际中更容易遇到码本崩溃问题，也就是大量码本向量得不到使用，模型只用其中一小部分词汇。官方实现里通常采用EMA指数滑动平均更新码本，每次将输入进来的特征向量做加权平均，让每个码本向量缓慢逼近它负责的输入特征。这样能明显提高码本利用率，训练也更稳定。

损失函数也不是单一的。整体由四个部分组成：像素重建损失、感知损失、对抗损失以及承诺损失。像素重建损失用L1或L2度量还原质量；感知损失通过预训练的VGG网络提取特征，比较真实图和重建图在特征空间的距离；对抗损失由PatchGAN判别器提供，让重建图更锐利自然；承诺损失则是约束编码器输出的特征不要偏离码本太远。

2.2 第二阶段的自回归Transformer

第一阶段训练好之后，图像就能被转成离散token序列了。做法很简单：把编码器和量化层推理一次，记录每个位置命中的码本索引，按光栅顺序展开，就得到了一个token序列。到这里，图像生成问题被彻底转化成了语言建模问题。

第二阶段用的是类似GPT的自回归Transformer。模型读取已知的token序列，逐位置预测下一个token。训练时用交叉熵损失，让模型输出的概率分布尽量贴近真实token。条件信息，比如类别标签或文本嵌入，会在每个时间步被拼接到输入序列中，让生成结果受控。

为什么第二阶段不直接复用第一阶段的效果？因为第一阶段只负责重建，并不具备从随机噪声或文本生成图像的能力。Transformer学到的才是“一张猫的图片的token序列长什么样”。生成时给一个统一起始符，模型一个token一个token地推下去，直到生成完整序列，再交给VQGAN解码器还原成像素图。

很多初学者容易把Transformer当成“图像生成器”，其实它产出的只是离散索引。真正把索引变成高清图像的是第一阶段的解码器G。这种解耦设计让两部分可以分别优化，第一阶段专注视觉重建质量，第二阶段专注序列分布建模，是工程上很漂亮的分工。

3. 环境搭建与PyTorch配置，避开版本陷阱

3.1 Anaconda创建独立环境

我强烈建议所有步骤都在独立的conda环境里完成，不要直接装在base环境，否则后面依赖冲突会非常痛苦。命令行操作如下：

conda create -n vqgan python=3.8 conda activate vqgan

选择Python 3.8不是随便选的。早期很多相关代码库对3.9和3.10兼容性都不好，尤其是pytorch-lightning和omegaconf的版本组合，3.8下踩坑最少。如果你用的是新版本PyTorch和改造过的代码仓库，Python 3.10也不是不行，但既然目标是复现流程，稳定优先。

3.2 GPU版PyTorch安装细节

PyTorch的安装是第一个容易出问题的地方。核心原则是保证CUDA驱动、显卡驱动和PyTorch三方兼容。先查看显卡驱动支持的CUDA版本：

nvidia-smi

输出右上角会显示Driver Version和CUDA Version，这个CUDA Version表示驱动最高支持的版本，不是当前激活的版本。只要你的PyTorch要求CUDA版本低于或等于驱动支持的最高版本，一般就能正常运行。比如你看到CUDA 12.1，安装cu118或cu121版本的PyTorch都没问题。

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

安装完成后，务必做一次GPU可用性验证：

import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))

这里我见过不少人翻车：torch.cuda.is_available()返回False，但nvidia-smi明明能看到显卡。原因几乎都是PyTorch和驱动不匹配，要么是装成了CPU版，要么是CUDA版本过高低版驱动不支持。重新卸载，对齐版本安装即可。

3.3 项目依赖与pytorch-lightning版本

VQGAN官方代码仓库用的是pytorch-lightning做训练框架。这个库版本更新很快，API改动也大。官方taming-transformers仓库在2021年开发，比较稳定的搭配是pytorch-lightning 1.x早期版本加omegaconf 2.0。如果你直接pip install最新版pytorch-lightning，大概率会因为API改名而报错。

建议按这样的顺序安装：

pip install omegaconf==2.0.0 pip install einops pip install transformers pip install torchmetrics==0.8.2 pip install pytorch-lightning==1.0.8

pytorch-lightning 1.0.8版本我对接官方仓库跑过多次，比较稳。然后安装taming-transformers本体：

pip install taming-transformers

如果你需要看或改源码，也可以直接git clone仓库并pip install -e .，这样后续调试时能直接在源码里打日志。我的建议是第二种，因为VQGAN的量化实现细节比较多，出问题时能直接看机制。

4. 官方仓库与预训练模型，先跑通再训练

4.1 仓库结构速览

taming-transformers仓库的核心内容分布在几个目录里。configs目录下是各种实验的YAML配置文件，训练和推理都会引用；taming目录下是模型实现代码，包括VQGAN模块、Transformer模块、损失函数和量化层；scripts目录下有训练和采样的入口脚本；data目录是数据集的下载说明。

配置文件不是摆设，它承担了模型参数、训练参数、数据路径等几乎所有配置。你打开一个YAML文件会发现里面有model、trainer、data三个大块，model块里又分ddconfig、lossconfig、n_embed、embed_dim等小节。理解了这个结构，你改模型参数和训练参数才会知道改哪里。

4.2 预训练模型怎么获取

官方在GitHub仓库里维护了一个预训练模型列表，覆盖了ImageNet、COCO、CelebA-HQ、S-FLCKR、FacesHQ等数据集。每个模型都对应一个配置文件，比如ImageNet 16384码本模型对应configs/imagenet/vqgan_imagenet_f16_16384.yaml。

权重文件通常放在release页面或Google Drive上。下载之后，你需要把权重路径和配置文件里ckpt_path字段对齐，或者直接把权重文件放到项目目录下的指定文件夹，然后在加载代码里写对路径。

我记得第一次跑的时候只下载了权重，忘记下载对应的YAML，结果维度不匹配直接报错，提醒一句，预训练权重和使用它的配置文件必须是一对一匹配的。不同码本大小、不同分辨率训练出来的模型结构完全不同，不能混用。

4.3 快速验证预训练模型

拿到权重后，先做一次图像重建验证，比直接训练靠谱得多。官方提供了reconstruction_demo.py脚本，原理是读取一张真实图片，经VQGAN编码再解码，返回重建结果。如果重建图像保留的主要特征都在，说明模型加载成功。

验证时注意输入尺寸要和配置文件的resolution字段保持一致。很多预训练模型是按256x256训练的，直接喂512x512图片会让编码器的下采样比例和码本数量对不上。正确的做法是先把图片resize到配置要求的分辨率。这也是新手最容易忽略的地方。

5. 从文本到高清图像：完整实操流程

5.1 两条路线怎么选

这里要先说明，纯官方的VQGAN代码里，Transformer阶段主要支持类别标签条件，没有直接封装“输入一句英文，输出一张图”的CLI命令。想实现真正的文本到图像，主流有两条路线。

第一条是自己写一个条件Transformer，把文本编码成向量，拼接到Transformer的条件位置里。这条路线最贴合VQGAN原理，适合学习和二次开发。第二条是直接用DALL·E复现类项目，它们底层就是VQGAN加文本条件Transformer，已经封装好了完整接口。对于只想快速出图验证效果的人，第二条更高效。

我会先讲如何用官方代码训练VQGAN和类条件Transformer，再给出一个文本条件注入的实现思路，并把两部分串起来完成从文本到高清图像的完整流程。

5.2 训练自己的VQGAN：数据准备与命令

VQGAN训练的第一步是准备数据。官方支持lmdb和文件夹两种数据加载方式，对新手来说，文件夹方式最友好。把训练图片放到一个目录下，然后在配置文件的data部分指向这个路径：

data: train: target: main.DataModuleFromConfig params: batch_size: 4 num_workers: 4 train: target: taming.data.custom.CustomTrain params: training_images_list_file: data/train_images.txt size: 256

size字段是关键，它表示输入图像的统一尺寸。建议直接用256，因为官方ImageNet模型的训练分辨率就是256。显存不够可以降到128，但码本压缩率是固定的，最终生成图像分辨率会相应降低。

生成训练集清单文件的方式很简单：

find /path/to/your/images -name "*.jpg" > data/train_images.txt

然后修改配置里的模型结构。比较敏感的是n_embed和embed_dim。n_embed代表码本大小，直接影响图像语义词汇量的丰富程度。官方在ImageNet上用了16384，初学者可以先用1024或4096跑通流程，再逐步增加。embed_dim是每个token的向量维度，官方常用256。这两个参数一改，第二阶段Transformer的输入维度也要跟着改。

启动训练的入口是train_vqgan.py：

python scripts/train_vqgan.py --config configs/custom_vqgan.yaml --gpus 0

训练过程中，你会看到输出很多损失数值。重点关注recon、perceptual和commit_loss。recon损失下降说明重建在变好，perceptual损失下降说明感知质量在提升，commit_loss如果波动剧烈，说明量化过程不稳定，可以适当调大β参数。

从我的实测经验来看，数据集比较小的时候，宁可batch_size小一些也不要强行拉大。VQGAN对batch size其实不算特别敏感，但判别器部分如果批次太小，训练容易震荡。建议batch_size不小于4，同时把图片分辨率设为256，这样显存占用可控。

5.3 训练类条件Transformer

第一阶段训练收敛后，你需要先用它把所有训练图片转换成token序列，然后再训练Transformer。官方提供了Indices2Keypoints或类似的预处理方式，但通用做法是写一个简单的批处理脚本，遍历图片，把编码后的码本索引保存为numpy数组。

Transformer训练中使用自回归任务，输入是一段token序列，输出是后移一位的预测结果。核心配置如下：

model: target: taming.models.vqgan.Transformer params: vocab_size: 4096 n_layer: 12 n_head: 8 n_embd: 512 cond_stage_key: class_label

vocab_size必须和VQGAN的n_embed一致，否则embedding查表维度对不上。n_layer和n_head决定模型容量，官方ImageNet版用24层16头，小数据集可以减半。训练命令类似：

python scripts/train_transformer.py --config configs/custom_transformer.yaml --gpus 0

训练时也要监控交叉熵损失，如果损失持续不降，优先怀疑配置里vocab_size和码本大小不一致，或者数据序列化过程出错。

5.4 把文本条件接进来，真正实现文本到图像

类条件生成的限制是你只能从固定集合里选类别，想自由输入文本就必须把条件编码换成文本向量。这里我分享一个稳定可行的做法，用CLIP的文本编码器抽取语义向量，然后当作条件注入到Transformer里。

文本编码部分可以用现成的transformers库加载CLIP模型：

from transformers import CLIPTextModel, CLIPTokenizer tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") text = "a red fox jumping over a log" inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=77) text_emb = text_encoder(**inputs).last_hidden_state # [1, 77, 512]

得到文本特征是77个token的嵌入向量。最简单有效的注入方式是取均值池化得到一个全局向量，再投影到与Transformer嵌入维度一致，拼在图像token序列的最前面，当作一个特殊的“指令token”。

如果你想让文本对每个图像token都有更强的控制力，可以把77个文本token的前若干个直接拼接进输入序列，让自回归模型在预测任意位置时都能注意到文本信息。这种做法更接近DALL·E的思路，但序列长度会变长，显存消耗也更大。

Transformer部分给一个最小可运行的模型示例框架：

import torch from torch import nn class TextCondTransformer(nn.Module): def __init__(self, vocab_size, cond_dim, n_layer, n_head, n_embd, max_seq_len): super().__init__() self.token_emb = nn.Embedding(vocab_size, n_embd) self.pos_emb = nn.Parameter(torch.zeros(1, max_seq_len, n_embd)) self.cond_proj = nn.Linear(cond_dim, n_embd) decoder_layer = nn.TransformerDecoderLayer(d_model=n_embd, nhead=n_head, batch_first=True) self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_layer) self.lm_head = nn.Linear(n_embd, vocab_size) def forward(self, tokens, text_cond): cond = self.cond_proj(text_cond).unsqueeze(1) x = self.token_emb(tokens) + self.pos_emb[:, :tokens.size(1)] x = self.decoder(x, cond) return self.lm_head(x)

训练时，text_cond就是从CLIP取到的全局文本向量，tokens是图像码本索引。推理时，输入起始符和文本条件，逐步采样下一个token，直到序列生成结束。最后把索引序列传入VQGAN解码器，就得到高清图像。

5.5 推理参数与生成质量调优

推理阶段有几个参数对图像质量影响非常大。第一个是温度系数，温度越低，采样越确定，但图像容易趋同；温度太高，结构容易乱。我的经验是从1.0开始调，往下降到0.7左右是一个比较好的平衡点。第二个是top-k采样，只从概率最高的k个token里采样，能有效避免低概率崩溃，通常设在100到300之间。

采样脚本核心逻辑类似：

with torch.no_grad(): for step in range(max_seq_len): logits = model(seq, text_cond) logits = logits[:, -1, :] / temperature logits = top_k_filter(logits, k=top_k) probs = torch.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) seq = torch.cat([seq, next_token], dim=1)

生成的token序列经过解码器后，图像分辨率取决于码本大小和编码器下采样比例。官方f16模型会对16x16的二维特征进行量化，对应256x256的输入图像，所以解码出来大约是256x256。想要更高清，可以在解码器后接一个超分模型，或者直接训练一个下采样比例更小的VQGAN。

6. 常见问题与排查技巧实录

6.1 高频报错速查表

6.2 训练不收敛，先查数据再调参

训练不收敛时，我的习惯是先不看损失曲线，先做一轮数据可视化。把训练集图片统一尺寸后直接喂给VQGAN做个重建测试，如果重建都糊，说明数据预处理或编码器配置有问题，跟Transformer无关。数据正常了，再检查两阶段之间token序列是否对齐。

还有一个容易踩的坑是数据集中图片长宽比差异过大。VQGAN默认按正方形处理，宽图会被强行拉伸变形。正确做法是先把图片中心裁剪成正方形或保持比例resize到短边等于目标尺寸再居中裁剪。这个预处理细节决定了感知损失能不能有效下降。

6.3 显存不够的三种实用策略

显存不够是最常见的硬件问题。第一种策略是降低batch_size到1，这是最直接的，但会带来BN统计不稳定，VQGAN本身没有太多BN，一般没问题。第二种策略是在配置文件中开启梯度累积，等效提升batch size。第三种策略是用mixed precision，pytorch-lightning里在Trainer中加precision=16参数，能省近一半显存，代价是损失数值看起来不太平滑，属于正常现象。

如果这三种都试了还爆显存，那只能从模型尺寸方面找突破口，比如减少Transformer的层数、降低embed_dim，或者把生成分辨率降到128。别指望一张4GB显卡能跑224层Transformer，结构上就注定了容量上限。

6.4 生成图像出现伪影和重复纹理

生成图像有伪影，通常有几个原因。第一是码本利用率过低，很多区域用的是同一个token，导致纹理重复。解决方法是加大码本，或者用更温和的码本初始化方式。第二是Transformer预测的token在局部区域自相冲突，比如左半边预测是绿色森林，右半边预测是树干，拼在一起就会出伪影。这个可以尝试降低采样温度，让全局一致性更强。

重复纹理还要考虑是不是训练数据里同类图像太多，模型产生了模式坍缩。这时候需要扩充数据多样性，而不是动模型参数。VQGAN这套管线对数据多样性的要求比普通GAN更高，因为Transformer部分本质上是语言模型，数据越多样越不容易学偏。

写在最后，再分享一点实操心得

从VQGAN开始学起，你会发现自己对图像生成模型的理解会上一个台阶。它不像现在的扩散模型那样直接端到端，而是强行把视觉信号拆分成离散token再重组，这个过程逼着你去理解表征学习、自回归建模和感知损失之间的关系。我建议动手复现一遍第一阶段，不一定要训练完整数据集，哪怕是几百张图的小数据集，也能直观感受到码本学习是怎么回事。

如果后续要扩展，可以试试两个方向：一是把Transformer换成扩散模型，在VQGAN的离散token序列上做扩散，这就是扩散语言模型的雏形；二是把VQGAN的码本直接迁移到多模态模型里，作为统一视觉词汇表。这套思路至今在很多前沿模型里依然能看到影子。你能从一篇论文、一份代码里吃透它，后续再学其他生成模型都会轻松很多。