Diffusion代码思路

还是有一些地方没搞懂

import numpy as np
import math
import torch.nn as nn
import torch
import torch.nn.functional as F
import time


class SinousoidalPosEmb(nn.Module): #位置编码网络
    def __init__(self,dim):
        super(SinousoidalPosEmb,self).__init__()
        self.dim = dim

    def forward(self,x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


class MLP(nn.Module):  #去噪网络
    def __init__(self, state_dim,action_dim,hidden_dim,device,time_dim=16):
        super(MLP, self).__init__()
        """
        state_dim: 
        action_dim: 
        hidden_dim: 
        device: device
        time_dim: 时间维度，每一步扩散的时间需要被编码
        """
        self.time_dim = time_dim
        self.device = device
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.hidden_dim = hidden_dim

        self.time_mlp = nn.Sequential(  #位置编码
            SinousoidalPosEmb(self.time_dim),
            nn.Linear(self.time_dim, self.time_dim*2),
            nn.Mish(),
            nn.Linear(self.time_dim*2, self.time_dim),
        )

        self.input_dim = self.state_dim + self.action_dim + self.time_dim
        self.mid_layer = nn.Sequential(
            nn.Linear(self.input_dim, self.hidden_dim),
            nn.Mish(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            nn.Mish(),
            nn.Linear(self.hidden_dim, self.hidden_dim),
            nn.Mish(),
        )
        self.final_layer = nn.Linear(self.hidden_dim, self.action_dim)

        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x, time, state):
        time_emb = self.time_mlp(time)
        x = torch.cat((x, state,time_emb), dim=1)
        x = self.mid_layer(x)
        return self.final_layer(x)


class WeightedLoss(nn.Module): #加权损失 采用继承
    def __init__(self):
        super(WeightedLoss, self).__init__()

    def forward(self, pred, target, weighted=1.0):
        loss = self._loss(pred, target)
        WeightedLoss = (loss* weighted).mean()
        return WeightedLoss

class WeightedL1(WeightedLoss):
    def _loss(self, pred, target):
        return torch.abs(pred - target)

class WeightedL2(WeightedLoss):
    def _loss(self, pred, target):
        return F.mse_loss(pred, target,reduction='none')
Losses = {
    'L1': WeightedL1,
    'L2': WeightedL2,
}

def extract(a,t,x_shape):  #取对于t的累乘α数据
    """ Extract values from a 1-D alpha tensor and reshape them to desired shape """
    b,*_ = x_shape
    out = a.gather(-1,t)
    return out.reshape(b,*((1,) * (len(x_shape)-1)))


class Diffusion(nn.Module): #扩散网络
    def __init__(self, loss_type,beta_schedule ='linear',clip_denoised = True,predict_epsilon=True,**kwargs):
        super(Diffusion, self).__init__()
        self.state_dim = kwargs["obs_dim"]
        self.action_dim = kwargs['act_dim']
        self.hidden_dim = kwargs['hidden_dim']
        self.T = kwargs['T'] #反传几步
        self.loss_type = loss_type
        self.clip_denoised = clip_denoised
        self.predict_epsilon = predict_epsilon
        self.device = torch.device(kwargs['device'])
        self.model = MLP(self.state_dim, self.action_dim, self.hidden_dim, self.device).to(kwargs['device'])

        if beta_schedule == 'linear':  #公式里的β间距划分
            betes = torch.linspace(1e-4, 0.02, self.T,dtype=torch.float32)

        alphas = 1.0 - betes   # α=1-β
        alphas_cumprod = torch.cumprod(alphas,axis=0) #αt的累乘 [1,2,3] -> [1,2,6]
        alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) # αt-1的累乘 [1,2,6] -> [1,1,2]


        #我们可能希望模型中的某些参数参数不更新（从开始到结束均保持不变），但又希望参数保存下来,就会用到 register_buffer()
        self.register_buffer("betes", betes)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)

        #前向过程
        # 根号累乘到αt
        self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        # 根号1-累乘到αt
        self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))

        #反向过程
        #后验过程的方差
        posterior_variance = betes * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        self.register_buffer("posterior_variance", posterior_variance)
        self.register_buffer("posterior_log_variance_clipped", torch.log(posterior_variance.clamp(min=1e-20)))
        #后验过程的均值
        self.register_buffer("sqrt_recip_alphas_cumprod",torch.sqrt(1.0/alphas_cumprod))
        self.register_buffer("sqrt_recipm1_alphas_cumprod",torch.sqrt(1.0/alphas_cumprod - 1))

        self.register_buffer("posterior_mean_coef1", betes * torch.sqrt(alphas_cumprod_prev)/(1.0 - alphas_cumprod))
        self.register_buffer("posterior_mean_coef2",(1.0-alphas_cumprod_prev)*torch.sqrt(alphas)/(1.0 - alphas_cumprod))

        self.loss_fn = Losses[loss_type]() #选择损失函数


    def q_posterior(self,x_start,x,t): #求后验概率
        posterior_mean = (extract(self.posterior_mean_coef1,t,x.shape) * x_start
        + extract(self.posterior_mean_coef2,t,x.shape) * x
        )
        posterior_variance = extract(self.posterior_variance,t,x.shape)
        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped,t,x.shape)
        return posterior_mean,posterior_variance,posterior_log_variance_clipped

    def predict_start_from_noise(self,x,t,pred_noise):
        return (extract(self.sqrt_recip_alphas_cumprod,t,x.shape) * x
                -extract(self.sqrt_recipm1_alphas_cumprod,t,x.shape)*pred_noise)


    def p_mean_variance(self,x,t,s):
        pred_noise = self.model(x,t,s)
        x_recon = self.predict_start_from_noise(x,t,pred_noise)
        x_recon.clamp_(-1,1) #下划线代表在内部操作，不用再赋值
        model_mean, posterior_variance,posterior_log_variance = self.q_posterior(x_recon,x,t)
        return model_mean, posterior_log_variance   #用log比较稳定


    def p_sample(self,x,t,s):
        b, *_,device = *x.shape,x.device
        model_mean,model_log_variance = self.p_mean_variance(x,t,s) #模型预测的均值和方差
        noise=torch.randn_like(x)

        nonzero_mask = (1 - (t == 0).float()).reshape(b,*((1,)*(len(x.shape)-1))) #T=0不加噪
        return model_mean + nonzero_mask * (0.5*model_log_variance).exp() * noise

    def p_sample_loop(self, state, shape, *args,**kwargs): #反向采样过程
        device = self.device
        batch_size = state.shape[0]
        x = torch.randn(shape, device=device,requires_grad=False)  #高斯噪声

        for i in reversed(range(0,self.T)):
            t = torch.full((batch_size,),i,device=device,dtype=torch.long)
            x = self.p_sample(x,t,state)

        return x


    def sample(self,state,*args,**kwargs):
        """
        state: [batch_size, state_dim]
        """
        batch_size = state.shape[0]
        shape = [batch_size,self.action_dim]  #初始化噪声的形状
        action = self.p_sample_loop(state,shape, *args,**kwargs) #ddpm不断扩散的过程
        return action.clamp_(-1,1) #防止action超过能和环境交互的边界

    # --------------------------------------------training------------------------------------------#

    def q_sample(self,x_start,t,noise):
        sample =(
            extract(self.sqrt_alphas_cumprod,t,x_start.shape) * x_start
            + extract(self.sqrt_one_minus_alphas_cumprod,t,x_start.shape) * noise
        )
        return sample
    def p_losses(self,x_start,state,t,weights=1.0):
        noise = torch.randn_like(x_start)
        x_noisy = self.q_sample(x_start,t,noise)
        x_recon = self.model(x_noisy,t,state)

        loss = self.loss_fn(x_recon,noise,weights)
        return loss


    def loss(self,x,state,weights=1.0):
        batch_size = len(x)
        t = torch.randint(0,self.T,(batch_size,),device=self.device).long()
        return self.p_losses(x,state,t,weights)

    def forward(self,state,*args,**kwargs): #采样过程
        return self.sample(state, *args,**kwargs)




if __name__ == '__main__':
    device = "cpu"
    x = torch.randn(256,2).to(device)
    state = torch.randn(256,11).to(device)
    model = Diffusion(loss_type="L2",obs_dim=11,act_dim=2,hidden_dim=256,T=100,device=device)
    action = model(state)

    loss =model.loss(x,state)
    print(f"action:{action};loss:{loss.item()}")