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

forward(input)

Mapping function for the MaskedLinear layer

Parameters:

input (Tensor) – input tensor to be transformed

Returns:

output tensor from the layer

Return type:

Tensor

classmethod from_pretrained(linear_module)

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

Parameters:

linear_module (nn.Linear) – Target layer to be transformed.

Returns:

New copy of the provided module with masked layers inserted to enable FedPM

Return type:

MaskedLinear