MaskedLinear#

class torch_uncertainty.layers.MaskedLinear(in_features, out_features, num_estimators, scale, bias=True, device=None, dtype=None)[source]#

Masksembles-style Linear layer.

This layer computes fully-connected operation for a given number of estimators (num_estimators) with a given scale.

Parameters:

  • in_features (int) – Number of input features of the linear layer.

  • out_features (int) – Number of channels produced by the linear layer.

  • num_estimators (int) – The number of estimators grouped in the layer.

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

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

  • groups – Number of blocked connections from input channels to output channels. Defaults to 1.

  • 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 MaskedLinear.

References

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