starling.samplers.ddim_sampler.DDIMSampler

class DDIMSampler

Bases: Module

Methods

__init__	An efficient sampler that generates samples 10x to 100x faster than the DDPM model.
add_module	Add a child module to the current module.
apply	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers	Return an iterator over module buffers.
children	Return an iterator over immediate children modules.
compile	Compile this Module's forward using torch.compile().
cpu	Move all model parameters and buffers to the CPU.
cuda	Move all model parameters and buffers to the GPU.
double	Casts all floating point parameters and buffers to double datatype.
eval	Set the module in evaluation mode.
extra_repr	Return the extra representation of the module.
float	Casts all floating point parameters and buffers to float datatype.
forward	Define the computation performed at every call.
generate_labels	Generate labels to condition the generative process on.
get_buffer	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state	Return any extra state to include in the module's state_dict.
get_parameter	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule	Return the submodule given by target if it exists, otherwise throw an error.
get_x_prev_and_pred_x0	Remove the noise from the latent iteratively.
half	Casts all floating point parameters and buffers to half datatype.
ipu	Move all model parameters and buffers to the IPU.
load_state_dict	Copy parameters and buffers from state_dict into this module and its descendants.
modules	Return an iterator over all modules in the network.
mtia	Move all model parameters and buffers to the MTIA.
named_buffers	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
p_sample	Take one step in the generative process.
parameters	Return an iterator over module parameters.
register_backward_hook	Register a backward hook on the module.
register_buffer	Add a buffer to the module.
register_forward_hook	Register a forward hook on the module.
register_forward_pre_hook	Register a forward pre-hook on the module.
register_full_backward_hook	Register a backward hook on the module.
register_full_backward_pre_hook	Register a backward pre-hook on the module.
register_load_state_dict_post_hook	Register a post-hook to be run after module's load_state_dict() is called.
register_load_state_dict_pre_hook	Register a pre-hook to be run before module's load_state_dict() is called.
register_module	Alias for add_module().
register_parameter	Add a parameter to the module.
register_state_dict_post_hook	Register a post-hook for the state_dict() method.
register_state_dict_pre_hook	Register a pre-hook for the state_dict() method.
requires_grad_	Change if autograd should record operations on parameters in this module.
sample	Sample the generative process using the DDIM model.
set_extra_state	Set extra state contained in the loaded state_dict.
set_submodule	Set the submodule given by target if it exists, otherwise throw an error.
share_memory	See torch.Tensor.share_memory_().
state_dict	Return a dictionary containing references to the whole state of the module.
to	Move and/or cast the parameters and buffers.
to_empty	Move the parameters and buffers to the specified device without copying storage.
train	Set the module in training mode.
type	Casts all parameters and buffers to dst_type.
xpu	Move all model parameters and buffers to the XPU.
zero_grad	Reset gradients of all model parameters.

Attributes

T_destination
call_super_init
dump_patches
training

__init__(ddpm_model, encoder_model, n_steps: int, ionic_strength: float = 150, ddim_discretize: str = 'uniform', ddim_eta: float = 0.0)

An efficient sampler that generates samples 10x to 100x faster than the DDPM model. Denoising diffusion implicit models (DDIM) do not require sampling the entire diffusion process to generate samples. The forward process is parameterized using non-Markovian diffusion processes, leading to short generative Markov chains that can be simulated in fewer steps.

References

[1] Ho, J., Jaini, P., Hariharan, B., Abbeel, P., & Duan, Y. (2020). Denoising diffusion implicit models. arXiv preprint arXiv:2012.02142.

Parameters:

• ddpm_model (_type_) – The trained DDPM model.

• n_steps (int) – The number of steps to simulate the generative process, smaller than the number of steps used to train the DDPM model.

• ddim_discretize (str, optional) – The discretization method for the generative process, by default “uniform”.

• ddim_eta (float, optional) – The noise level for the generative process, a number between 0.0 and 1.0. 0.0 adds no noise to the generative process, 1.0 adds the maximum noise. This number interpolates between deterministic and stochastic generative processes, by default 0.0.

Raises:

NotImplementedError – If the discretization method is not implemented.

generate_labels(labels: str) → Tensor

Generate labels to condition the generative process on.

Parameters:

labels (str) – A sequence to generate labels from.

Returns:

The labels to condition the generative process on.

Return type:

torch.Tensor

sample(num_conformations: int, labels: Tensor, repeat_noise: bool = False, temperature: float = 1.0, show_per_step_progress_bar: bool = True, batch_count: int = 1, max_batch_count: int = 1, constraint=None) → Tensor

Sample the generative process using the DDIM model.

Parameters:

• num_conformations (int) – Number of conformations to generate.

• labels (torch.Tensor) – The labels to condition the generative process on.

• repeat_noise (bool, optional) – _description_, by default False

• temperature (float, optional) – _description_, by default 1.0

• show_per_step_progress_bar (bool, optional) – whether to show progress bar per step.

• batch_count (int, optional) – The batch count for the progress bar, by default 1

• max_batch_count (int, optional) – The maximum batch count for the progress bar, by default 1

Returns:

The generated distance maps.

Return type:

torch.Tensor

p_sample(x: Tensor, c: Tensor, t: Tensor, attention_mask: Tensor, step: int, index: int, repeat_noise: bool = False, temperature: float = 1.0) → Tuple[Tensor, Tensor, Tensor]

Take one step in the generative process.

Parameters:

• x (torch.Tensor) – The tensor to remove noise from.

• c (torch.Tensor) – The labels to condition the generative process on.

• t (torch.Tensor) – The timestep to sample the generative process at.

• step (int)

• index (int) – _description_

• repeat_noise (bool, optional) – _description_, by default False

• temperature (float, optional) – _description_, by default 1.0

Returns:

_description_

Return type:

_type_

get_x_prev_and_pred_x0(predicted_noise: Tensor, index: int, x: Tensor, temperature: float, repeat_noise: bool) → Tuple[Tensor, Tensor]

Remove the noise from the latent iteratively.

Parameters:

• predicted_noise (torch.Tensor) – The noise predicted by the DDPM model.

• index (int) – The index of the timestep.

• x (torch.Tensor) – The latent to remove the noise from.

• temperature (float) – The temperature to use for the generative process

• repeat_noise (bool) – Whether to repeat the noise

Returns:

The previous latent.

Return type:

Tuple[torch.Tensor, torch.Tensor]