විසරණසම්භාවිතාව ආකෘති නිරූපණය කිරීම (DDPM)

163from typing import Tuple, Optional
164
165import torch
166import torch.nn.functional as F
167import torch.utils.data
168from torch import nn
169
170from labml_nn.diffusion.ddpm.utils import gather

173class DenoiseDiffusion:

178    def __init__(self, eps_model: nn.Module, n_steps: int, device: torch.device):

184        super().__init__()
185        self.eps_model = eps_model

188        self.beta = torch.linspace(0.0001, 0.02, n_steps).to(device)

191        self.alpha = 1. - self.beta

193        self.alpha_bar = torch.cumprod(self.alpha, dim=0)

195        self.n_steps = n_steps

197        self.sigma2 = self.beta

199    def q_xt_x0(self, x0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

209        mean = gather(self.alpha_bar, t) ** 0.5 * x0

211        var = 1 - gather(self.alpha_bar, t)

213        return mean, var

215    def q_sample(self, x0: torch.Tensor, t: torch.Tensor, eps: Optional[torch.Tensor] = None):

225        if eps is None:
226            eps = torch.randn_like(x0)

229        mean, var = self.q_xt_x0(x0, t)

231        return mean + (var ** 0.5) * eps

233    def p_sample(self, xt: torch.Tensor, t: torch.Tensor):

247        eps_theta = self.eps_model(xt, t)

249        alpha_bar = gather(self.alpha_bar, t)

251        alpha = gather(self.alpha, t)

253        eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5

256        mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)

258        var = gather(self.sigma2, t)

261        eps = torch.randn(xt.shape, device=xt.device)

263        return mean + (var ** .5) * eps

265    def loss(self, x0: torch.Tensor, noise: Optional[torch.Tensor] = None):

274        batch_size = x0.shape[0]

276        t = torch.randint(0, self.n_steps, (batch_size,), device=x0.device, dtype=torch.long)

279        if noise is None:
280            noise = torch.randn_like(x0)

283        xt = self.q_sample(x0, t, eps=noise)

285        eps_theta = self.eps_model(xt, t)

288        return F.mse_loss(noise, eps_theta)