Skip to content

Module bastionlab.torch.psg.convert


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.