fl4health.model_bases.masked_layers.masked_linear module

class MaskedLinear(in_features, out_features, bias=True, device=None, dtype=None)

Bases: Linear

__init__(in_features, out_features, bias=True, device=None, dtype=None)

Implementation of masked linear layers.

Like regular linear layers (i.e., nn.Linear module), a masked linear layer has a weight and a bias. However, the weight and the bias do not receive gradient in back propagation. Instead, two score tensors - one for the weight and another for the bias - are maintained. In the forward pass, the score tensors are transformed by the Sigmoid function into probability scores, which are then used to produce binary masks via bernoulli sampling. Finally, the binary masks are applied to the weight and the bias. During training, gradients with respect to the score tensors are computed and used to update the score tensors.

Note: the scores are not assumed to be bounded between 0 and 1.

Parameters:

• in_features (int) – size of each input sample

• out_features (int) – size of each output sample

• bias (bool) – If set to False, the layer will not learn an additive bias. Default: True

weight

weights of the module.

bias

bias of the module.

weight_score

learnable scores for the weights. Has the same shape as weight.

bias_score

learnable scores for the bias. Has the same shape as bias.

forward(input)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

classmethod from_pretrained(linear_module)

Return an instance of MaskedLinear whose weight and bias have the same values as those of linear_module.

Return type:

MaskedLinear