所有之前：这篇用来记录一些经典算法的代码实现，只有标题没有内容的是TODO，这里只会记录领域的最经典的算法，针对经典算法的变体和改进算法在其他文章。
Transformer
下面是一个最基本的Transformer实现，主要包含以下几个部分：
- 基本模块
- 多头注意力模块
- Token和位置嵌入模块
- LayerNorm
- 前馈神经网络FFN
- $EncoderLayer -> Encoder$
- $DecoderLayer -> Decoder$
- $Encoder+Decoder -> Transformer$
多头注意力MHA
多头注意力机制的公式如下：
$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O $$其中，每一个头的计算分别为：
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V $$对于图像数据和序列数据，多头注意力机制的实现是不同，但是核心原理是相同的。对于序列数据，每个位置对应一个token；对于图像数据，每个位置对应一个patch块或者一个像素块，如果是$patch$块，需要展开到一维向量。这里实现的为序列数据的多头注意力机制,初始$shape$为$(batch,time,d_model)$
import torch
from torch import nn
import torch.functional as F
import math
class MultiHeadAttention(nn.Module):
def __init__(self,d_model,n_head):
super(MultiHeadAttention, self).__init__()
assert d_model % n_head == 0
self.d_model = d_model
self.n_head = n_head
self.w_k = nn.Linear(d_model,d_model)
self.w_q = nn.Linear(d_model,d_model)
self.w_v = nn.Linear(d_model,d_model)
self.w_conbine = nn.Linear(d_model,d_model,bias=False) # 对应W^O
self.softmax = nn.Softmax(dim=-1)
def forward(self,q,k,v,mask=None):
batch,time,dimension = q.shape() # q,k,v shape: (batch,time,d_model)
n_d = self.d_model // self.n_head
q,k,v = self.w_q(q),self.w_k(k),self.w_v(v)
print(q.shape, batch, time, self.n_head, n_d)
q = q.view(batch,time,self.n_head,n_d).transpose(1,2) # self.n_head * n_d = d_model
k = k.view(batch,time,self.n_head,n_d).transpose(1,2) # shape: (batch,n_head,time,n_d)
v = v.view(batch,time,self.n_head,n_d).transpose(1,2)
# compute attention by the formula
score = q @ k.transpose(-2,-1) / math.sqrt(n_d) # (batch,n_head,time,time)
if mask is not None:
score = score.masked_fill(mask==0,float('-inf'))
attn = self.softmax(score) @ v # (batch,n_head,time,n_d)
attn = attn.transpose(1,2).contiguous().view(batch,time,dimension) # (batch,time,dimension)
attn = self.w_conbine(attn) # (batch,time,dimension)
return attn
d_model = 512
n_head = 8
mha = MultiHeadAttention(d_model,n_head)
x = torch.randn(2,10,d_model)
y = mha(x,x,x) # self-multi-head-attention
print(f"shape: {y.shape}\n输出:{y}") # torch.Size([2, 10, 512])
Token和位置嵌入模块
Token Embedding 是将离散的词（token）转换为连续的向量表示的过程。通常通过查找表（如 nn.Embedding）将每个词的索引映射为一个高维向量，使模型能够处理和学习词之间的语义关系。 公式表示：
$$ \text{TokenEmbedding}(x) = E[x] $$其中 $E$ 是嵌入矩阵，$x$ 是词的索引。TokenEmbedding是一个专门的领域，这里不做深究，直接使用torch自己的类。
Position Embedding为模型提供序列的顺序信息，Transformer 结构本身不具备序列顺序信息，因此需要引入 Position Embedding，将每个位置的信息编码为向量，并与 Token Embedding 相加。常见做法有两种：
- 可学习的位置嵌入（Learnable Position Embedding）：每个位置对应一个可训练的向量。
- 固定的位置嵌入（Sinusoidal Position Embedding）：使用三角函数编码序列位置信息。
以固定的位置嵌入为例,公式如下:
$$ PE_{(pos,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$$$ PE_{(pos,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$import torch
from torch import nn
import torch.functional as F
import math
class TokenEmbedding(nn.Embedding):
def __init__(self,vocab_size,d_model):
super(TokenEmbedding,self).__init__(vocab_size,d_model,padding_idx=1)
#----------------------------------------------------------------------------
class PositionalEncoding(nn.Module):
def __init__(self,d_model,device,max_len=5000):
super(PositionalEncoding,self).__init__()
self.encoding = torch.zeros(max_len,d_model,device=device) # (max_len,d_model)
self.encoding.requires_grad = False # 固定的位置嵌入不需要梯度
pos = torch.arange(0,max_len,device=device).float().unsqueeze(1) # (max_len)-> (max_len,1)
_2i = torch.arange(0,d_model,2,device=device).float() # (d_model/2,)
# PE(pos,2i) = sin(pos/10000^(2i/d_model))
self.encoding[:,0::2] = torch.sin(pos / (10000 ** (_2i/d_model)))
# PE(pos,2i+1) = cos(pos/10000^(2i/d_model))
self.encoding[:,1::2] = torch.cos(pos / (10000 ** (_2i/d_model)))
def forward(self,x):
batch,seq_len,dimension = x.size()
# return x + self.encoding[:seq_len,:] # 直接返回带位置嵌入的x,多余的位置嵌入不需要
return self.encoding[:seq_len,:] # 只返回位置嵌入
#----------------------------------------------------------------------------
class TransformerEmbedding(nn.Module):
def __init__(self,vocab_size,d_model,device,max_len=5000,dropout_p=0.1):
super(TransformerEmbedding,self).__init__()
self.token_embedding = TokenEmbedding(vocab_size,d_model)
self.position_embedding = PositionalEncoding(d_model,device,max_len)
self.dropout = nn.Dropout(dropout_p)
# def forward(self,x):
# x = self.token_embedding(x)
# x = self.position_embedding(x)
# x = self.dropout(x)
# return x
def forward(self,x):
t_emb = self.token_embedding(x)
p_emb = self.position_embedding(t_emb)
return self.dropout(t_emb + p_emb)
LayerNorm
LayerNorm（层归一化）是一种常用的归一化方法，主要用于深度学习模型中提升训练稳定性和加速收敛。它的作用是对每个样本的特征维度进行归一化，使其均值为0、方差为1，并引入可学习的缩放和偏移参数。
公式如下：
$$ \text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta $$$x$：输入向量 $\mu$：均值，$\mu = \frac{1}{N} \sum_{i=1}^N x_i$ $\sigma^2$：方差，$\sigma^2 = \frac{1}{N} \sum_{i=1}^N (x_i - \mu)^2$ $\gamma$ 和 $\beta$：可学习的缩放和偏移参数 $\epsilon$：防止除零的小常数
Tip:LayerNorm（层归一化）和 BatchNorm（批归一化）主要区别如下： BatchNorm 是对每个特征维度在整个 batch 上做归一化，依赖于 batch 的统计信息。 LayerNorm 是对每个样本的所有特征维度做归一化，和 batch 大小无关，对于序列数据(batch,time,d)来说，就是对每个time的d维数据进行归一化；对于图像数据来说，就是在图像每个通道的像素点进行归一化。
import torch
from torch import nn
import torch.functional as F
import math
class LayerNorm(nn.Module):
def __init__(self,d_model,eps=1e-9):
super(LayerNorm,self).__init__()
# two learnable parameters
self.gamma = nn.Parameter(torch.ones(d_model))
self.beta = nn.Parameter(torch.zeros(d_model))
# one small constant to prevent division by zero
self.eps = eps
def forward(self,x):
mean = x.mean(-1,keepdim=True) # (batch,time,1)
std = x.std(-1,keepdim=True) # (batch,time,1)
return self.gamma * (x - mean) / (std + self.eps) + self.beta # by the formula
x = torch.randn(2,10,512)
ln = LayerNorm(512)
y = ln(x)
print(f"shape: {y.shape}\n输出:{y}")
FFN前馈网络（别看名字，实际很简单）
FFN 的结构如下：
- 包含两个线性变换（全连接层）和一个激活函数（通常是ReLU或GELU）。
- 对每个时间步/位置单独处理，不进行序列间的信息交互。
公式表示：
$$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$import torch
from torch import nn
import torch.functional as F
import math
class ffn(nn.Module):
def __init__(self,d_model,d_hidden,dropout=0.1):
super(ffn,self).__init__()
self.fc1 = nn.Linear(d_model,d_hidden)
self.fc2 = nn.Linear(d_hidden,d_model)
self.dropout = nn.Dropout(dropout) # dropout放在激活函数后，减少激活的神经元
self.relu = nn.ReLU()
def forward(self,x):
return self.fc2(self.dropout(self.relu(self.fc1(x)))) # by the formula
集成所有模块，构建完整的Transformer模型
首先构建EncoderLayer和DecoderLayer，然后将它们堆叠成Encoder和Decoder。
import torch
from torch import nn
import torch.functional as F
import math
class EncoderLayer(nn.Module):
def __init__(self,d_model,n_head,d_hidden,dropout_p=0.1):
super(EncoderLayer,self).__init__()
self.mha = MultiHeadAttention(d_model,n_head)
self.ln1 = LayerNorm(d_model)
self.dropout = nn.Dropout(dropout_p)
self.ffn = ffn(d_model,d_hidden,dropout_p)
self.ln2 = LayerNorm(d_model)
self.dropout2 = nn.Dropout(dropout_p)
def forward(self,x,mask=None):
_x = x
x = self.mha(x,x,x,mask)
x = self.dropout(x)
x = self.ln1(x + _x)
_x = x
x = self.ffn(x)
x = self.dropout2(x)
x = self.ln2(x + _x)
return x
class DecoderLayer(nn.Module):
def __init__(self,d_model,n_head,d_hidden,dropout_p=0.1):
super(DecoderLayer,self).__init__()
self.s_mha = MultiHeadAttention(d_model,n_head) # self-multi-head-attention with mask
self.ln1 = LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout_p)
self.cross_mha = MultiHeadAttention(d_model,n_head) # cross-multi-head-attention
self.ln2 = LayerNorm(d_model)
self.dropout2 = nn.Dropout(dropout_p)
self.ffn = ffn(d_model,d_hidden,dropout_p)
self.ln3 = LayerNorm(d_model)
self.dropout3 = nn.Dropout(dropout_p)
def forward(self,dec,enc,time_mask=None,src_mask=None):
_x = dec
x = self.s_mha(dec,dec,dec,time_mask) # 避免解码器看到未来的信息,下三角掩码
x = self.dropout1(x)
x = self.ln1(x + _x)
if enc is not None:
_x = x
x = self.cross_mha(x,enc,enc,src_mask) # 对于padding的mask
x = self.dropout2(x)
x = self.ln2(x + _x)
_x = x
x = self.ffn(x)
x = self.dropout3(x)
x = self.ln3(x + _x)
return x
首先使用之前写好的组件构成Encoder和Decoder,然后完成整个Transformer，Mask是Transformer中最关键的一部分,怎么构建两个掩码以及两个掩码的作用分别是什么
class Encoder(nn.Module):
def __init__(self,vocab_size,d_model,n_head,d_hidden,device,max_len=5000,dropout_p=0.1,num_layers=6):
super(Encoder,self).__init__()
self.embedding = TransformerEmbedding(vocab_size,d_model,device,max_len,dropout_p)
self.layers = nn.ModuleList(
[EncoderLayer(d_model,n_head,d_hidden,dropout_p) for _ in range(num_layers)]
)
def forward(self,x,src_mask=None):
x = self.embedding(x)
for layer in self.layers:
x = layer(x,src_mask)
return x
class Decoder(nn.Module):
def __init__(self,vocab_size,d_model,n_head,d_hidden,device,max_len=5000,dropout_p=0.1,num_layers=6):
super(Decoder,self).__init__()
self.embedding = TransformerEmbedding(vocab_size,d_model,device,max_len,dropout_p)
self.layers = nn.ModuleList(
[DecoderLayer(d_model,n_head,d_hidden,dropout_p) for _ in range(num_layers)]
)
self.fc_out = nn.Linear(d_model,vocab_size,bias=False)
def forward(self,dec,enc,time_mask=None,src_mask=None):
x = self.embedding(dec)
for layer in self.layers:
x = layer(x,enc,time_mask,src_mask)
x = self.fc_out(x)
return x
class Transformer(nn.Module):
def __init__(self,src_pad_idx,trg_pad_idx,enc_voc_size,dec_voc_size,max_len,d_model,n_head,d_hidden,n_layers,dropout_p,device):
super(Transformer,self).__init__()
self.encoder = Encoder(enc_voc_size,d_model,n_head,d_hidden,device,max_len,dropout_p,n_layers)
self.decoder = Decoder(dec_voc_size,d_model,n_head,d_hidden,device,max_len,dropout_p,n_layers)
self.src_pad_idx = src_pad_idx
self.trg_pad_idx = trg_pad_idx
self.device = device
def make_casual_mask(self,q,k):
'''
构建下三角掩码（time mask）
'''
len_q,len_k = q.size(1),k.size(1)
mask = torch.tril(torch.ones((len_q,len_k),device=self.device)).type(torch.BoolTensor).to(self.device)
return mask
def make_pad_mask(self,q,k,pad_idx_q,pad_idx_k):
'''
构建padding掩码（src mask）
SHAPE: (batch,time,len_q,len_k)
'''
len_q,len_k = q.size(1),k.size(1)
q = q.ne(pad_idx_q).unsqueeze(1).unsqueeze(3)
q = q.repeat(1,1,1,len_k) # (batch,1,len_q,len_k)
k = k.ne(pad_idx_k).unsqueeze(1).unsqueeze(2)
k = k.repeat(1,1,len_q,1) # (batch,1,len_q,len_k)
mask = q & k
return mask
def forward(self,src,trg):
src_mask = self.make_pad_mask(src,src,self.src_pad_idx,self.src_pad_idx)
trg_mask = self.make_pad_mask(trg,trg,self.trg_pad_idx,self.trg_pad_idx) * self.make_casual_mask(trg,trg)
src_trg_mask = self.make_pad_mask(trg,src,self.trg_pad_idx,self.src_pad_idx)
enc = self.encoder(src,src_mask)
out = self.decoder(trg,enc,trg_mask,src_trg_mask)
return out
if __name__ == "__main__":
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
src_pad_idx = 1
trg_pad_idx = 1
enc_voc_size = 11
dec_voc_size = 11
max_len = 20
d_model = 512
n_head = 8
d_hidden = 2048
n_layers = 6
dropout_p = 0.1
model = Transformer(src_pad_idx,trg_pad_idx,enc_voc_size,dec_voc_size,max_len,d_model,n_head,d_hidden,n_layers,dropout_p,device).to(device)
src = torch.tensor([[2,5,6,4,3,9,7,8,1,1,1],[3,4,5,6,7,8,9,1,1,1,1]],dtype=torch.long).to(device) # (batch,time)
trg = torch.tensor([[2,7,4,3,5,9,8,1,1,1,1,1],[3,5,6,7,8,9,1,1,1,1,1,1]],dtype=torch.long).to(device) # (batch,time)
out = model(src,trg[:,:-1])
print(f"shape: {out.shape}\n输出:{out}")
Diffusion Models
原理
具体的一些公式推导请看李宏毅老师的视频： 李宏毅Diffusion
基本概念
（以图像生成为例）首先Diffusion的基本运作过程是，首先Sample出一个全是噪声的图片，然后经过一步一步的Denoise最终得到图片，Denoise的步数是事先确定好的。
在Diffusion中，深度学习模型的唯一作用就是预测噪声，其他过程均由数学公式推理得到
进一步说，不同情况的噪声需要的处理肯定是不一样的，一个Denoise模型肯定无法做到处理不同的情况的噪声，因此Denoise模型需要一个“噪声严重程度”作为输入，在Diffusion中称为时间步。
Denoise内部的工作是，接受带Noise的图片，对噪声进行预测，然后使用带Noise的图片减去预测噪声那么就可以得到。
我们清楚Noise predictor模型的输入输出之后，该怎么训练这个模型呢？已知模型的输入-输出对是Noise图片加上“噪声程度”-Noise数据对，我们需要人为创造这个数据对出来，具体来说，从image list中拿一张图片，然后从GasuionN中采样出一张纯噪声的图片，加了一定的步数之后，就得到了很多数据对（Noise图片、时间步–噪声），这个过程称为Diffusion Process。
现在，已经可以从一个纯噪声得到一张图片，但是我们不能根据噪声随机生成图片，所以我们需要把文字考虑进来，让模型考虑进文字之后再预测噪声。
完整算法
首先是Training的过程，从immage list中取出一张干净的图，之后从1-T中随机取一个值作为时间步，然后从Gaussion中取样出一个noise，之后就是一步一步的梯度更新训练模型。这种方法和我们的想法有些不一样，我们想像中noise是一步一步加进去的，但是在实际中是一下子直接加进去的。
具体的加噪公式推导过程如下：
然后是推理过程，也就是Sampling这个过程。首先，采样得到一个纯噪声图片X_T，然后进行T次公式中的循环，$\epsilon_\theta$表示预测噪声的模型。
但是为什么要加这个z噪声：每次都要重新采样出一个噪声，但是只有当t>1的时候.这个z的使用和生成语句中一个问题类似：为什么每次总要SAMPLE而不是直接取Mean（概率最大）？在The curious case of nerual text degeneration论文中有一个分析，当直接取Mean时候，会一直重复相同的回答和出现跳帧现象，并且他们分析得到人在写文章的时候并不会一定会选概率最大的词。
简单代码实现
Diffusion已经有了许多变体，下面是一个简单的实现，变体不会在这篇里更新。
import torch
from torch import nn
import math
import torch.nn.functional as F
import numpy as np
import time
class WeightedLoss(nn.Module):
def __init__(self):
super(WeightedLoss,self).__init__()
def forward(self,pred,target,weighted=1.0):
loss = self._loss(pred,target)
weighted_loss = (loss * weighted).mean() # mean by batch
return weighted_loss
class L1Loss(WeightedLoss):
def _loss(self,pred,target):
return torch.abs(pred - target)
class L2Loss(WeightedLoss):
def _loss(self,pred,target):
return (pred - target) ** 2
Losses = {
'l1':L1Loss,
'l2':L2Loss
}
class SinusoidalPosEmb(nn.Module):
def __init__(self,dim):
super(SinusoidalPosEmb,self).__init__()
self.dim = dim
def forward(self,t):
'''
扩散模型时间步的位置编码公式和Transformer中的位置编码公式是类似的。
'''
device = t.device
half_dim = self.dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim,device=device) * -emb)
emb = t[:,None] * emb[None,:]
emb = torch.cat((emb.sin(),emb.cos()),dim=-1)
return emb
def extract(a,t,x_shape):
'''
a: (T,) t: (B,) x_shape: (B,...)
'''
batch_size = t.shape[0]
out = a.gather(-1,t) # (B,)
return out.view(batch_size, *((1,) * (len(x_shape) - 1))) # (B,1,1,1...)
class MLP(nn.Module):
'''
这是一个简单的多层感知机(MLP)模型，结合了时间步的位置编码，用于处理状态和动作的输入（强化学习）
'''
def __init__(self,state_dim,action_dim,hidden_dim,device,t_dim):
super(MLP,self).__init__()
self.device = device
self.t_dim = t_dim
self.a_tim = action_dim
self.time_mlp = nn.Sequential(
SinusoidalPosEmb(t_dim),
nn.Linear(t_dim,t_dim*2),
nn.Mish(),
nn.Linear(t_dim*2,t_dim)
)
input_dim = state_dim + action_dim + t_dim
self.mid_layer = nn.Sequential(
nn.Linear(input_dim,hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim,hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim,hidden_dim),
nn.Mish(),
)
self.f_layer = nn.Linear(hidden_dim,action_dim)
def init_weights(self):
'''
初始化模型权重,好的权重初始化可以帮助模型更快收敛，提升训练效果。
'''
for m in self.modules():
if isinstance(m,nn.Linear):
nn.init.xavier_normal_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
def forward(self,state,action,t):
t_emb = self.time_mlp(t) # (batch,t_dim)
x = torch.cat((state,action,t_emb),dim=-1) # (batch,state_dim+action_dim+t_dim)
x = self.mid_layer(x) # (batch,hidden_dim)
action_delta = self.f_layer(x) # (batch,action_dim)
return action_delta
class Diffusion(nn.Module):
def __init__(self,loss_type,beta_schedule='linear',clip_denoise=True,predict_epsilon=True,**kwargs):
'''
loss_type: 损失函数类型
beta_schedule: beta(时间)的离散方式，常见的有'linear'（线性）和'cosine'（余弦）。
clip_denoise: 是否在去噪过程中裁剪输出，以防止数值过大。
'''
super(Diffusion,self).__init__()
self.device = kwargs.get('device','cpu')
self.t_dim = kwargs.get('t_dim',16)
self.state_dim = kwargs.get('state_dim',10)
self.action_dim = kwargs.get('action_dim',2)
self.hidden_dim = kwargs.get('hidden_dim',256)
self.loss_type = loss_type
self.beta_schedule = beta_schedule
self.clip_denoise = clip_denoise
self.T = kwargs.get('timesteps',1000)
self.model = MLP(self.state_dim,self.action_dim,self.hidden_dim,self.device,self.t_dim).to(self.device)
self.kwargs = kwargs
self.model.init_weights()
if beta_schedule == 'linear':
betas = torch.linspace(1e-4,0.02,self.kwargs.get('timesteps', 1000),dtype=torch.float32)
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas,dim=0) # [1,2,3] -> [1,1*2,1*2*3]
alphas_cumprod_prev = torch.cat((torch.tensor([1.],dtype=torch.float32),alphas_cumprod[:-1]),dim=0)
self.register_buffer('betas',betas)
self.register_buffer('alphas',alphas)
self.register_buffer('alphas_cumprod',alphas_cumprod)
self.register_buffer('alphas_cumprod_prev',alphas_cumprod_prev)
self.register_buffer('sqrt_alphas_cumprod',torch.sqrt(alphas_cumprod))
# （前向过程）
self.register_buffer('sqrt_alphas_cumprob',torch.sqrt(alphas_cumprod))
self.register_buffer('sqrt_one_minus_alphas_cumprod',torch.sqrt(1. - alphas_cumprod))
# （反向过程）
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
self.register_buffer('posterior_variance',posterior_variance)
# 用于从xt求得x0
self.register_buffer('sqrt_recip_alphas_cumprod',torch.sqrt(1. / alphas_cumprod))
self.register_buffer('sqrt_recipm1_alphas_cumprod',torch.sqrt(1. / alphas_cumprod - 1))
self.register_buffer('posterior_mean_coef1',betas *torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
self.register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
self.loss_func = Losses[loss_type]()
def forward(self,state,*args,**kwargs):
return self.sample(state,*args,**kwargs)
def sample(self,state,*args,**kwargs):
batch_size = state.shape[0]
shape = [batch_size,self.action_dim]
action = self.p_sample_loop(state,shape,*args,**kwargs)
return action.clamp(-1.,1.)
def q_posterior(self,x0,xt,t):
'''
计算后验分布q(x_{t-1}|x_t,x_0)的均值和方差
'''
posterior_mean = (extract(self.posterior_mean_coef1,t,x0.shape) * x0 +
extract(self.posterior_mean_coef2,t,x0.shape) * xt)
posterior_variance = extract(self.posterior_variance,t,x0.shape)
posterior_log_variance = torch.log(posterior_variance.clamp(min=1e-20))
return posterior_mean,posterior_variance,posterior_log_variance
def p_sample_loop(self,state,shape,*args,**kwargs):
device = self.kwargs.get('device','cpu')
batch_size = shape[0]
x = torch.randn(shape,device=device,requires_grad=False) # 标准DDPM噪声不需要梯度
for i in reversed(range(0,self.kwargs.get('timesteps',1000))):
t = torch.full((batch_size,),i,dtype=torch.long,device=device)
x = self.p_sample(x,t,state)
return x
def predict_x0_from_noise(self,x,t,noise):
'''
根据公式计算x0
'''
return (extract(self.sqrt_recip_alphas_cumprod,t,x.shape) * x -
extract(self.sqrt_recipm1_alphas_cumprod,t,x.shape) * noise)
def p_mean_variance(self,x,t,state):
pred_noise = self.model(state,x,t)
x_0 = self.predict_x0_from_noise(x,t,pred_noise)
x_0.clamp_(-1.,1.) if self.clip_denoise else x_0
model_mean,posterior_variance,posterior_log_variance = self.q_posterior(x_0,x,t)
return model_mean,posterior_log_variance
def p_sample(self,x,t,state):
'''
x: 当前的噪声状态xt
t: 当前的时间步t
state: 状态
'''
model_mean,model_log_variance = self.p_mean_variance(x,t,state)
noise = torch.randn_like(x) if t.sum() > 0 else 0. # 如果t=0,则不添加噪声
return model_mean + torch.exp(0.5 * model_log_variance) * noise
def q_sample(self,x0,t,noise):
sample = (extract(self.sqrt_alphas_cumprod,t,x0.shape) * x0 +
extract(self.sqrt_one_minus_alphas_cumprod,t,x0.shape) * noise)
return sample
def p_losses(self,x0,state,t,weights):
noise = torch.randn_like(x0)
x_noisy = self.q_sample(x0,t,noise)
x_recon = self.model(state,x_noisy,t)
loss = self.loss_func(x_recon,noise,weights)
return loss
def loss(self,x,state,weights=1.0):
batch_size = x.shape[0]
t = torch.randint(0,self.T,(batch_size,),device=self.device).long()
return self.p_losses(x,state,t,weights)
if __name__ == '__main__':
x = torch.randn(256,2)
state = torch.randn(256,11)
model = Diffusion(loss_type='l2',beta_schedule='linear',clip_denoise=True,predict_epsilon=True,
device='cpu',t_dim=16,state_dim=11,action_dim=2,hidden_dim=256,timesteps=100)
result = model(state)
print(result)
loss = model.loss(x,state)
print(f"loss: {loss.item()}")
VAE（变分自编码器）
VAE也是一个非常经典的生成模型，虽然现在主流已经是Diffusion模型，但是它的原理和思想还是非常有启发性的，尤其是对于理解diffusion中的一些概念（如潜在空间、重参数化技巧等）非常有帮助。