MaskedConv2d#

class torch_uncertainty.layers.MaskedConv2d(in_channels, out_channels, kernel_size, num_estimators, scale, stride=1, padding=0, dilation=1, groups=1, bias=True, device=None, dtype=None)[source]#

Masksembles-style Conv2d layer.

Parameters:

  • in_channels (int) – Number of channels in the input image.

  • out_channels (int) – Number of channels produced by the convolution.

  • kernel_size (Union[int, tuple[int, int]]) – Size of the convolving kernel.

  • num_estimators (int) – Number of estimators in the ensemble.

  • scale (float) – The scale parameter for the masks.

  • stride (Union[int, tuple[int, int]]) – Stride of the convolution. Defaults to 1.

  • padding (int, tuple or str) – Padding added to all four sides of the input. Defaults to 0.

  • dilation (Union[int, tuple[int, int]]) – Spacing between kernel elements. Defaults to 1.

  • groups (int) – Number of blocked connexions from input channels to output channels for each estimator. Defaults to 1.

  • bias (bool) – If True, adds a learnable bias to the output. Defaults to True.

  • device (Any | None) – The desired device of returned tensor. Defaults to None.

  • dtype (Any | None) – The desired data type of returned tensor. Defaults to None.

Warning

Be sure to apply a repeat on the batch at the start of the training if you use MaskedConv2d.

References

[1] Masksembles for Uncertainty Estimation, Nikita Durasov, Timur Bagautdinov, Pierre Baque, Pascal Fua.