Timesteps#

class vision_architectures.utils.timesteps.TimestepSampler(total_timesteps, strategy, gamma_alpha=1.0, gamma_spread=0.8, signal_to_noise_ratios=None, temperature=2.0)[source]#

Bases: Module

Sampel timesteps using a strategy.

Parameters:

  • total_timesteps (int) – Total number of timesteps.

  • strategy (Literal['uniform', 'gamma', 'importance']) – Sampling strategy.

  • gamma_alpha (float) – Used when strategy == gamma. Alpha parameter for the gamma distribution.

  • gamma_spread (float) – Used when strategy == gamma. Controls how spread out the distribution is.

  • signal_to_noise_ratios (Optional[Tensor]) – Used when strategy == importance. Signal-to-noise ratios of the noise schedule.

  • temperature (float) – Used when strategy == importance. Temperature for the softmax distribution.

__init__(total_timesteps, strategy, gamma_alpha=1.0, gamma_spread=0.8, signal_to_noise_ratios=None, temperature=2.0)[source]#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(num_timesteps)[source]#

Sample timesteps based on the specified strategy.

Parameters:

num_timesteps (int) – Number of timesteps to sample.