Stochastic

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.

`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:

Name	Type	Description	Default
`p`	`Normal`	The prior generative distribution p(z_i\|z_{i+1}) (or p(z_L)).	required
`p_params`	`Tensor`	The parameters of the prior generative distribution.	required
`q`	`Normal`	The inference distribution q(z_i\|z_{i+1}) (or q(z_L\|x)).	required
`q_params`	`Tensor`	The parameters of the inference distribution.	required
`mode_pred`	`bool`	Whether the model is in prediction mode.	required
`analytical_kl`	`bool`	Whether to compute the KL divergence analytically or using Monte Carlo estimation.	required
`z`	`Tensor`	The sampled latent tensor.	required

`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:

Name	Type	Description	Default
`p_params`	`Tensor`	The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}).	required
`q_params`	`Union[Tensor, 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`.	`None`
`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. Default is `None`.	`None`
`force_constant_output`	`bool`	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`.	`False`
`analytical_kl`	`bool`	Whether to compute the KL divergence analytically or using Monte Carlo estimation. Default is `False`.	`False`
`mode_pred`	`bool`	Whether the model is in prediction mode. Default is `False`.	`False`
`use_uncond_mode`	`bool`	Whether to use the uncoditional distribution p(z) to sample latents in prediction mode. Default is `False`.	`False`
`var_clip_max`	`Union[float, None]`	The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. Default is `None`.	`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:

Name	Type	Description	Default
`sampling_distrib`	`Normal`	The Gaussian distribution from which latent tensor is sampled.	required
`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.	required
`mode_pred`	`bool`	Whether the model is prediction mode.	required
`use_uncond_mode`	`bool`	Whether to use the uncoditional distribution p(z) to sample latents in prediction mode.	required

`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:

Name	Type	Description	Default
`p_params`	`Tensor`	The input tensor to be processed.	required
`var_clip_max`	`float`	The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.	required

`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:

Name	Type	Description	Default
`p_params`		The input tensor to be processed.	required
`var_clip_max`	`float`	The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.	required

`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:

Name	Type	Description	Default
`q_params`	`Tensor`	The input tensor to be processed.	required
`var_clip_max`	`float`	The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped.	required