annotated_deep_learning_paper_implementations 中的生成式AI:从 GPT 到 Stable Diffusion
生成式AI技术正以前所未有的速度改变着我们与数字内容交互的方式。在开源项目annotated_deep_learning_paper_implementations中,我们可以找到从GPT到Stable Diffusion等一系列里程碑式生成模型的实现代码和详细注释。本文将带你深入探索这些模型的工作原理、实现细节以及如何在实际场景中应用它们。
GPT模型家族:语言生成的突破
GPT(Generative Pre-trained Transformer)系列模型标志着自然语言处理领域的重大突破。项目中提供了基于GPT-2架构的实现,包括完整的模型结构和多种优化技术。
GPT-2模型结构解析
GPT-2模型采用了Transformer架构的 decoder-only 设计,主要由以下组件构成:
- 12层Transformer blocks,每层包含一个多注意力头模块和一个前馈神经网络
- 12个注意力头,隐藏层维度为768
- 使用了正弦位置编码和残差连接
项目中的实现代码位于labml_nn/lora/gpt2目录下。以下是模型初始化的核心代码片段:
from labml_nn.lora.gpt2 import GPTModel
model: GPTModel
# 初始化GPT2模型
self.model = GPTModel(
vocab_size=self.vocab_size,
d_model=self.d_model,
n_heads=self.n_heads,
n_layers=self.n_layers,
n_positions=self.n_positions,
layer_norm_epsilon=self.layer_norm_epsilon,
lora_r=self.lora_r
)
LoRA微调技术
项目特别关注了参数高效微调技术LoRA(Low-Rank Adaptation)在GPT-2上的应用。LoRA通过冻结预训练模型参数,仅训练少量适配器参数来实现模型微调,这显著降低了计算资源需求。
docs/zh/lora/experiment.html中详细介绍了使用LoRA微调GPT-2的完整流程。关键配置参数包括:
- LoRA秩(lora_r):32
- 学习率:1e-4
- 批大小:32
- 上下文长度:512
以下是加载预训练GPT-2模型并应用LoRA的代码示例:
# 从HuggingFace加载预训练GPT-2模型
hf_model = AutoModelForCausalLM.from_pretrained("gpt2")
state_dict = hf_model.state_dict()
# 权重映射和转换
mapping = {
'transformer.wte.weight': 'token_embedding.weight',
'transformer.wpe.weight': 'position_embedding.weight',
'transformer.ln_f.weight': 'final_norm.weight',
'transformer.ln_f.bias': 'final_norm.bias',
'lm_head.weight': 'lm_head.weight'
}
# 初始化带有LoRA的GPT2模型
self.model = GPTModel(
vocab_size=self.vocab_size,
d_model=self.d_model,
n_heads=self.n_heads,
n_layers=self.n_layers,
n_positions=self.n_positions,
layer_norm_epsilon=self.layer_norm_epsilon,
lora_r=self.lora_r
)
# 加载权重
self.model.load_state_dict(new_state_dict, strict=False)
文本生成采样策略
生成高质量文本不仅依赖于模型架构,还需要精心设计的采样策略。docs/sampling/experiment.html展示了如何在GPT-2模型上应用不同的采样技术,包括:
- 贪婪采样(Greedy Sampling)
- 温度采样(Temperature Sampling)
- Top-K采样
- Nucleus采样(Top-P采样)
这些技术直接影响生成文本的多样性和质量。以下是使用不同采样策略生成文本的示例代码:
from transformers import GPT2Tokenizer, GPT2LMHeadModel
def sample(model: GPT2LMHeadModel, tokenizer: GPT2Tokenizer, sampler: Sampler,
prompt: str, length: int = 20, num_samples: int = 1, temperature: float = 1.0):
# 对输入文本进行编码
input_ids = tokenizer.encode(prompt, return_tensors='pt').to(model.device)
# 设置生成参数
generation_kwargs = {
"max_length": input_ids.shape[1] + length,
"num_return_sequences": num_samples,
"do_sample": True,
"temperature": temperature,
"pad_token_id": tokenizer.eos_token_id,
}
# 根据采样器类型调整参数
if isinstance(sampler, TopKSampler):
generation_kwargs["top_k"] = sampler.k
elif isinstance(sampler, NucleusSampler):
generation_kwargs["top_p"] = sampler.p
generation_kwargs["top_k"] = 0 # 确保top_k不覆盖top_p
# 生成文本
output = model.generate(input_ids,** generation_kwargs)
# 解码并返回结果
return [tokenizer.decode(sequence, skip_special_tokens=True) for sequence in output]
Stable Diffusion:文本到图像的革命
Stable Diffusion是一种强大的文本到图像生成模型,它通过潜在空间中的扩散过程实现高质量图像生成。项目在docs/diffusion/stable_diffusion/index.html提供了完整实现,包括模型架构、采样算法和应用脚本。
潜在扩散模型架构
Stable Diffusion的核心是潜在扩散模型(Latent Diffusion Model),它由三个主要组件构成:
- 自动编码器(Autoencoder):将图像压缩到低维潜在空间,减少计算复杂度
- U-Net:在潜在空间执行扩散过程,包含自注意力机制
- CLIP文本编码器:将文本提示转换为模型可理解的嵌入向量

docs/diffusion/stable_diffusion/model/index.html详细介绍了各个组件的实现细节。U-Net结构是模型的核心,它通过一系列下采样和上采样操作,结合自注意力机制捕捉图像的全局特征:
class UNet(nn.Module):
def __init__(self,
in_channels: int = 4,
out_channels: int = 4,
channels: List[int] = [320, 640, 1280, 1280],
attention_levels: List[bool] = [False, True, True, True],
n_heads: int = 8,
tf_layers: int = 1,
d_cond: int = 768):
super().__init__()
# 输入卷积层
self.input_blocks = nn.ModuleList([
TimestepEmbedSequential(Conv2d(in_channels, channels[0], 3, padding=1))
])
# 下采样块
self.down_blocks = nn.ModuleList()
for i in range(len(channels)):
# 添加注意力和残差块
# ...
# 添加下采样层(除最后一层外)
if i != len(channels) - 1:
self.down_blocks.append(TimestepEmbedSequential(
nn.Sequential(
nn.Upsample(scale_factor=0.5, mode="nearest"),
Conv2d(channels[i], channels[i+1], 3, padding=1)
)
))
# 中间块
self.middle_block = TimestepEmbedSequential(
# 添加中间层和注意力
# ...
)
# 上采样块
self.up_blocks = nn.ModuleList()
# ...
# 输出卷积层
self.out = nn.Sequential(
normalization(channels[0]),
SiLU(),
Conv2d(channels[0], out_channels, 3, padding=1)
)
文本引导的图像生成
Stable Diffusion的魔力在于它能够根据文本描述生成高度相关的图像。docs/diffusion/stable_diffusion/scripts/text_to_image.html提供了完整的文本到图像生成实现。
生成过程主要包括以下步骤:
- 文本编码:使用CLIP模型将文本提示转换为嵌入向量
- 潜在空间采样:在潜在空间中执行扩散过程,生成图像的潜在表示
- 解码:使用自动编码器将潜在表示转换为实际图像
以下是文本到图像生成的核心代码:
class Txt2Img:
model: LatentDiffusion
def __init__(self, *,
checkpoint_path: Path,
sampler_name: str,
n_steps: int = 50,
ddim_eta: float = 0.0,
):
# 加载预训练模型
self.model = load_model(checkpoint_path)
self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
self.model.to(self.device)
# 初始化采样器
if sampler_name == 'ddim':
self.sampler = DDIMSampler(self.model, n_steps=n_steps, ddim_eta=ddim_eta)
elif sampler_name == 'ddpm':
self.sampler = DDPMSampler(self.model)
@torch.no_grad()
def __call__(self, *,
dest_path: str,
batch_size: int = 3,
prompt: str,
h: int = 512, w: int = 512,
uncond_scale: float = 7.5,
):
# 设置参数
c = 4 # 通道数
f = 8 # 下采样因子
prompts = batch_size * [prompt]
with torch.cuda.amp.autocast():
# 获取文本嵌入
if uncond_scale != 1.0:
un_cond = self.model.get_text_conditioning(batch_size * [""])
else:
un_cond = None
cond = self.model.get_text_conditioning(prompts)
# 在潜在空间采样
x = self.sampler.sample(cond=cond,
shape=[batch_size, c, h // f, w // f],
uncond_scale=uncond_scale,
uncond_cond=un_cond)
# 解码生成图像
x = self.model.decode_first_stage(x)
# 保存图像
save_images(x, dest_path, 'txt2img')
采样算法
Stable Diffusion提供了多种采样算法来平衡生成质量和速度,主要包括:
- DDPM采样:原始扩散概率模型采样,需要较多步骤(通常1000步)
- DDIM采样:加速扩散隐式模型,可在50步内生成高质量图像
docs/diffusion/stable_diffusion/sampler/index.html详细比较了这些采样算法的性能。以下是DDIM采样器的核心实现:
class DDIMSampler:
def __init__(self, model: LatentDiffusion, n_steps: int = 50, ddim_eta: float = 0.0):
self.model = model
self.n_steps = n_steps
self.ddim_eta = ddim_eta
# 计算DDIM时间步
self.ddim_timesteps = np.asarray(list(range(0, 1000, 1000//n_steps)))
self.ddim_timesteps = np.append(self.ddim_timesteps, 999)
# 预计算alpha和sigma
self.ddim_sigmas = torch.sqrt((1 - self.ddim_alphas_prev) * self.ddim_eta**2)
@torch.no_grad()
def sample(self, cond, shape, uncond_scale=1.0, uncond_cond=None):
# 初始化随机噪声
x = torch.randn(shape, device=self.model.device)
# 反向扩散过程
for i in reversed(range(self.n_steps)):
# 计算时间步
timestep = self.ddim_timesteps[i]
ts = torch.full((x.shape[0],), timestep, device=x.device, dtype=torch.long)
# 双重条件(unconditional guidance)
if uncond_cond is not None and uncond_scale != 1.0:
# 混合条件和无条件输出
# ...
# 应用DDIM更新规则
# ...
return x
实际应用与扩展
annotated_deep_learning_paper_implementations项目不仅提供了基础模型实现,还包含多种创意应用和扩展功能:
图像到图像转换
除了文本到图像生成,项目还实现了图像到图像的转换功能。docs/diffusion/stable_diffusion/scripts/image_to_image.html展示了如何根据文本提示修改现有图像,同时保留原始图像的结构和风格。
图像修复
docs/diffusion/stable_diffusion/scripts/in_paint.html实现了图像修复功能,允许用户通过文本提示修改图像的特定区域。这在图像编辑和修复中非常实用。
Flash Attention加速
为提高生成速度,项目集成了Flash Attention技术到U-Net的注意力模块中。这可以在RTX A6000等GPU上实现近50%的性能提升,显著减少图像生成时间。
总结与展望
annotated_deep_learning_paper_implementations项目为生成式AI的学习和应用提供了宝贵资源。从GPT系列的文本生成到Stable Diffusion的图像创作,项目全面覆盖了当今最先进的生成模型。
通过详细的代码注释和实验教程,开发者可以深入理解这些复杂模型的工作原理,并快速将其应用到实际项目中。无论是研究人员还是AI爱好者,都能从中获得启发和实用知识。
随着生成式AI技术的不断发展,我们可以期待看到更多创新应用和优化方法的出现。这个项目将持续跟踪最新研究进展,为开源社区提供高质量的实现参考。
建议读者通过以下资源深入学习:
- 官方文档:docs/index.html
- GPT-2实现:labml_nn/lora/gpt2
- Stable Diffusion源码:labml_nn/diffusion/stable_diffusion
- 示例脚本:docs/diffusion/stable_diffusion/scripts
通过实践这些实现,你将能够构建自己的生成式AI应用,并参与到这场AI创作革命中来。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考



