nn¶
This module tries to mirror the torch.nn module, offering all available layers to be readily replaced by calling the convert_to_memory_saving() function.
This module contains the following members:
This module also contains the submodule memsave_torch.nn.functional, which tries to mirror the torch.nn.functional module, offering layer operations as functions.
- memsave_torch.nn.convert_to_memory_saving(model: Module, *, conv2d=True, conv1d=True, conv3d=True, batchnorm2d=True, relu=True, maxpool2d=True, dropout=True, verbose=False, clone_params=False) Module¶
Converts the given model to it’s MemSave version, with options to choose which layer types to replace.
The clone_params option should be used when you plan on using both models simultaneously. Otherwise, the grad accumulation for one model wll affect the other (since their weights are the same Tensor object). For an example, see tests/test_layers.py.
- Parameters:
model (nn.Module) – The input model
conv2d (bool, optional) – Whether to replace nn.Conv2d layers
conv1d (bool, optional) – Whether to replace nn.Conv1d layers
conv3d (bool, optional) – Whether to replace nn.Conv3d layers
batchnorm2d (bool, optional) – Whether to replace nn.BatchNorm2d layers
relu (bool, optional) – Whether to replace nn.ReLU layers
maxpool2d (bool, optional) – Whether to replace nn.MaxPool2d layers
dropout (bool, optional) – Whether to replace nn.Dropout layers
verbose (bool, optional) – Whether to print which layers were replaced
clone_params (bool, optional) – Whether to clone the layer parameters or use directly
- Returns:
The converted memory saving model
- Return type:
memsavemodel (nn.Module)
Learnable Layers¶
MemSaveConv1d. |
|
MemSaveConv2d. |
|
MemSaveConv3d. |
|
Differentiability-agnostic 1d transpose convolution layer. |
|
Differentiability-agnostic 2d transpose convolution layer. |
|
Differentiability-agnostic 3d transpose convolution layer. |
Activations and Pooling Layers¶
MemSaveReLU. |
|
MemSaveMaxPool2d. |
Normalization Layers¶
MemSaveBatchNorm2d. |