Diffusion原理及代码实现

Diffusion原理及代码实现

代码实现

我们采用手写数字MNIST数据集来验证，将原始维度为[batch,1,28,28]，将其展平为[batch,784]。

核心代码（基于Unet）

# 定义一个简单的 U-Net 结构的编码器-解码器网络用于扩散模型
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1),
            nn.Tanh(),
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

# 定义扩散模型的核心
class DiffusionModel(nn.Module):
    def __init__(self, timesteps=1000):
        super(DiffusionModel, self).__init__()
        self.timesteps = timesteps
        self.network = UNet()

    def forward(self, x):
        return self.network(x)

    def noise_schedule(self, t):
        """返回在时间步t处的噪声因子"""
        return torch.linspace(1e-4, 0.02, self.timesteps)[t]

    def forward_diffusion_sample(self, x0, t):
        """在给定时间步t将噪声加入图像x0"""
        #生成与 x0 尺寸相同的随机噪声张量，噪声值服从标准正态分布（均值为0，标准差为1）
        noise = torch.randn_like(x0)
        
        alpha = self.noise_schedule(t)