functional

This module tries to mirror the torch.nn.functional module, offering layer operations as functions, where you need to provide the inputs to the layer, the layer parameters and any other options it might need.

This module contains the following members:

memsave_torch.nn.functional.batch_normMemSave(input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05) Tensor

Functional form of the memory saving batch_norm.

Parameters:
  • input – Input to the network [B, C, H, W]

  • running_mean – running_mean

  • running_var – running_var

  • weight – weight

  • bias – bias

  • training – training

  • momentum – momentum

  • eps – eps

Returns:

Output of the network [B, C, H, W]

Return type:

torch.Tensor

memsave_torch.nn.functional.convMemSave(input, weight, bias, stride, padding, dilation, groups, transposed, output_padding) Tensor

Functional form of the memory saving convolution.

Input can be 3D, 4D or 5D.

Parameters:
  • input – input [B, C_in, (None, H, D), W]

  • weight – weight

  • bias – bias

  • stride – stride

  • padding – padding

  • dilation – dilation

  • groups – groups

  • transposed – transposed

  • output_padding – output_padding

No Longer Returned:

torch.Tensor: Output of the conv operation [B, C_out, H_out, W_out]

memsave_torch.nn.functional.dropoutMemSave(x, p, training) Tensor

Functional form of the memory saving dropout.

Parameters:
  • x – Input to the network

  • p – Probability of elements being zeroed

  • training – Whether the layer is in training mode (no dropout applied in eval)

Returns:

Output of the network

Return type:

torch.Tensor

memsave_torch.nn.functional.maxpool2dMemSave(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False) Tensor | Tuple[Tensor, Tensor]

Functional form of the memory saving max-pooling.

Parameters:
  • input – input [B, C, H, W]

  • kernel_size – kernel_size

  • stride – stride

  • padding – padding

  • dilation – dilation

  • ceil_mode – ceil_mode

  • return_indices – If true, also returns the max_indices

Returns:

Output of the maxpool operation [B, C, H_out, W_out] (and optionally the indices when return_indices=True)

Return type:

Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

memsave_torch.nn.functional.reluMemSave(x) Tensor

Functional form of the memory saving relu.

Parameters:

x – Input to the network

Returns:

Output of the network

Return type:

torch.Tensor