所有之前：CodeByMySelf 系列之 Diffusion Models 的实现记录，变体不会在这篇里更新。
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()}")