File size: 823 Bytes
dfd1909 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | from numpy import ndarray
from torch import Tensor
import numpy as np
class BetaSchedule:
@staticmethod
def linear(timesteps:int, start:float=1e-4, end:float =2e-2) -> ndarray:
"""
linear schedule
"""
return np.linspace(start, end, timesteps)
@staticmethod
def cosine(timesteps:int, s:float=0.008) -> ndarray:
"""
cosine schedule
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
"""
steps:int = timesteps + 1
x = np.linspace(0, steps, steps)
alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
betas:ndarray = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
return np.clip(betas, a_min=0, a_max=0.999)
|