torch_blue.vi.VIConv2d

class torch_blue.vi.VIConv2d(in_channels: int, out_channels: int, kernel_size: int | Tuple[int, int], stride: int | Tuple[int, int] = 1, padding: str | int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', variational_distribution: Distribution | List[Distribution] = MeanFieldNormal(), prior: Distribution | List[Distribution] = MeanFieldNormal(), rescale_prior: bool = False, kaiming_initialization: bool = True, prior_initialization: bool = False, return_log_probs: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None)

Bases: _VIConvNd

Applies a 2D convolution over an input signal composed of several input planes.

Equivalent of Conv2d with variational inference. See its documentation for usage.

In addition to those arguments, this class accepts VIkwargs.

This module’s random variables are

• (“weight”, “bias”) if bias == True

• (“weight”, ) if bias == False

Parameters:

• torch_args – The same arguments and keyword arguments as the pytorch version Conv2d (documentation here)

• VIkwargs – Several standard keyword arguments. See VIkwargs for details.

forward(input_: torch.Tensor) → torch.Tensor

Forward computation.

Applies a 2D convolution to a tensor of shape [N, C_in, H_in, W_in]. Where N is the batch size, C are channels, and [H, W] are height and width of the input.

Parameters:

input (Tensor) – Input tensor of shape [N, C_in, H_in, W_in].

Returns:

output – Output tensor of shape [N, C_out, H_out, W_out]. Auto-sampling will add a sample dimension at the start for the overall output.

Return type:

Tensor