fl4health.model_bases.masked_layers.masked_layers_utils module

convert_to_masked_model(original_model)

Given a model, convert every one of its layers to a masked layer of the same kind, if applicable.

Parameters:

original_model (nn.Module) – Module to be converted to a masked module

Returns:

New copy of the original model but with masked layers injected to enable FedPM

Return type:

nn.Module

is_masked_module(module)

Checks whether the provided module is a masked module of the kind supported.

Parameters:

module (nn.Module) – Module to be checked

Returns:

True if the module is a masked type and False otherwise.

Return type:

bool