cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
(timesteps, s=0.008)
| 165 | |
| 166 | |
| 167 | def cosine_beta_schedule(timesteps, s=0.008): |
| 168 | """ |
| 169 | cosine schedule |
| 170 | as proposed in https://openreview.net/forum?id=-NEXDKk8gZ |
| 171 | """ |
| 172 | steps = timesteps + 1 |
| 173 | x = np.linspace(0, steps, steps) |
| 174 | alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2 |
| 175 | alphas_cumprod = alphas_cumprod / alphas_cumprod[0] |
| 176 | betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) |
| 177 | return np.clip(betas, a_min=0, a_max=0.999) |
| 178 | |
| 179 | |
| 180 | class GaussianDiffusion(nn.Module): |