Module bastionlab.torch.psg.convert

Functions

expand_layer(layer: torch.nn.modules.module.Module, max_batch_size: int) ‑> Optional[torch.nn.modules.module.Module]

Returns an expanded version of given layer if supported, else None.

expand_weights(module: torch.nn.modules.module.Module, max_batch_size: int) ‑> None

Recursively converts the layers of a model to their expanded counterpart in bastionlab.torch.psg.nn.

Args: module: model whose weights must be expanded. max_batch_size: maximum size of the batches that will be processed by the model.

parent_layer(module, name: str) ‑> Tuple[torch.nn.modules.module.Module, str]

Returns the parent module of the layer corresponding to the given name (path) in the given module and the attribute name of the layer wrt its parent.