MemSaveBatchNorm2d¶
- class memsave_torch.nn.MemSaveBatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None)¶
MemSaveBatchNorm2d.
Inits a BatchNorm2d layer with the given params
- Parameters:
num_features – num_features
eps – eps
momentum – momentum
affine – affine
track_running_stats – track_running_stats
device – device
dtype – dtype
- forward(x)¶
Forward pass.
- Parameters:
x – Input to the network [B, C, H, W]
- Returns:
Output [B, C, H, W]
- Return type:
- classmethod from_nn_BatchNorm2d(bn2d: BatchNorm2d)¶
Converts a nn.BatchNorm2d layer to MemSaveBatchNorm2d.
- Parameters:
bn2d – The nn.BatchNorm2d layer
- Returns:
The MemSaveBatchNorm2d object
- Return type:
obj
Hint
The usage is the same as torch.nn.BatchNorm2d
For usage examples, please refer to the linked torch documentation