Skip to content

Stochastic

Source

Script containing the common basic blocks (nn.Module) reused by the LadderVAE architecture.

NormalStochasticBlock

Bases: Module

Stochastic block used in the Top-Down inference pass.

Algorithm: - map input parameters to q(z) and (optionally) p(z) via convolution - sample a latent tensor z ~ q(z) - feed z to convolution and return.

NOTE 1: If parameters for q are not given, sampling is done from p(z).

NOTE 2: The restricted KL divergence is obtained by first computing the element-wise KL divergence (i.e., the KL computed for each element of the latent tensors). Then, the restricted version is computed by summing over the channels and the spatial dimensions associated only to the portion of the latent tensor that is used for prediction.

__init__(c_in, c_vars, c_out, conv_dims=2, kernel=3, transform_p_params=True, vanilla_latent_hw=None, use_naive_exponential=False)

Parameters:

  • c_in (int) –

    The number of channels of the input tensor.

  • c_vars (int) –

    The number of channels of the latent space tensor.

  • c_out (int) –

    The output of the stochastic layer. Note that this is different from the sampled latent z.

  • conv_dims (int, default: 2 ) –

    The number of dimensions of the convolutional layers (2D or 3D). Default is 2.

  • kernel (int, default: 3 ) –

    The size of the kernel used in convolutional layers. Default is 3.

  • transform_p_params (bool, default: True ) –

    Whether a transformation should be applied to the p_params tensor. The transformation consists in a 2D convolution ()conv_in_p()) that maps the input to a larger number of channels. Default is True.

  • vanilla_latent_hw (int, default: None ) –

    The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL). Default is None.

  • use_naive_exponential (bool, default: False ) –

    If False, exponentials are computed according to the alternative definition provided by StableExponential class. This should improve numerical stability in the training process. Default is False.

compute_kl_metrics(p, p_params, q, q_params, mode_pred, analytical_kl, z)

Compute KL (analytical or MC estimate) and then process it, extracting composed versions of the metric. Specifically, the different versions of the KL loss terms are: - kl_elementwise: KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)]. - kl_samplewise: KL term associated to each sample in the batch [Shape: (batch, )]. - kl_samplewise_restricted: KL term only associated to the portion of the latent tensor that is used for prediction and summed over channel and spatial dimensions [Shape: (batch, )]. - kl_channelwise: KL term associated to each sample and each channel [Shape: (batch, ch, )]. - kl_spatial: KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)]

Parameters:

  • p (Normal) –

    The prior generative distribution p(z_i|z_{i+1}) (or p(z_L)).

  • p_params (Tensor) –

    The parameters of the prior generative distribution.

  • q (Normal) –

    The inference distribution q(z_i|z_{i+1}) (or q(z_L|x)).

  • q_params (Tensor) –

    The parameters of the inference distribution.

  • mode_pred (bool) –

    Whether the model is in prediction mode.

  • analytical_kl (bool) –

    Whether to compute the KL divergence analytically or using Monte Carlo estimation.

  • z (Tensor) –

    The sampled latent tensor.

forward(p_params, q_params=None, forced_latent=None, force_constant_output=False, analytical_kl=False, mode_pred=False, use_uncond_mode=False, var_clip_max=None)

Parameters:

  • p_params (Tensor) –

    The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}).

  • q_params (Union[Tensor, None], default: None ) –

    The tensor resulting from merging the bu_value tensor at the same hierarchical level from the bottom-up pass and the p_params tensor. Default is None.

  • forced_latent (Union[Tensor, None], default: None ) –

    A pre-defined latent tensor. If it is not None, than it is used as the actual latent tensor and, hence, sampling does not happen. Default is None.

  • force_constant_output (bool, default: False ) –

    Whether to copy the first sample (and rel. distrib parameters) over the whole batch. This is used when doing experiment from the prior - q is not used. Default is False.

  • analytical_kl (bool, default: False ) –

    Whether to compute the KL divergence analytically or using Monte Carlo estimation. Default is False.

  • mode_pred (bool, default: False ) –

    Whether the model is in prediction mode. Default is False.

  • use_uncond_mode (bool, default: False ) –

    Whether to use the uncoditional distribution p(z) to sample latents in prediction mode. Default is False.

  • var_clip_max (Union[float, None], default: None ) –

    The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. Default is None.

get_z(sampling_distrib, forced_latent, mode_pred, use_uncond_mode)

Sample a latent tensor from the given latent distribution.

Latent tensor can be obtained is several ways: - Sampled from the (Gaussian) latent distribution. - Taken as a pre-defined forced latent. - Taken as the mode (mean) of the latent distribution. - In prediction mode (mode_pred==True), can be either sample or taken as the distribution mode.

Parameters:

  • sampling_distrib (Normal) –

    The Gaussian distribution from which latent tensor is sampled.

  • forced_latent (Union[Tensor, None]) –

    A pre-defined latent tensor. If it is not None, than it is used as the actual latent tensor and, hence, sampling does not happen.

  • mode_pred (bool) –

    Whether the model is prediction mode.

  • use_uncond_mode (bool) –

    Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.

process_p_params(p_params, var_clip_max)

Process the input parameters to get the prior distribution p(z_i|z_{i+1}) (or p(z_L)).

Processing consists in: - (optionally) 2D convolution on the input tensor to increase number of channels. - split the resulting tensor into two chunks, the mean and the log-variance. - (optionally) clip the log-variance to an upper threshold. - define the normal distribution p(z) given the parameter tensors above.

Parameters:

  • p_params (Tensor) –

    The input tensor to be processed.

  • var_clip_max (float) –

    The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.

process_q_params(q_params, var_clip_max, allow_oddsizes=False)

Process the input parameters to get the inference distribution q(z_i|z_{i+1}) (or q(z|x)).

Processing consists in: - convolution on the input tensor to double the number of channels. - split the resulting tensor into 2 chunks, respectively mean and log-var. - (optionally) clip the log-variance to an upper threshold. - (optionally) crop the resulting tensors to ensure that the last spatial dimension is even. - define the normal distribution q(z) given the parameter tensors above.

Parameters:

  • p_params

    The input tensor to be processed.

  • var_clip_max (float) –

    The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.

sample_from_q(q_params, var_clip_max)

Given an input parameter tensor defining q(z), it processes it by calling process_q_params() method and sample a latent tensor from the resulting distribution.

Parameters:

  • q_params (Tensor) –

    The input tensor to be processed.

  • var_clip_max (float) –

    The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.