torchscan.modules¶
The modules subpackage contains tools for inspection of modules.
FLOPs¶
Related to the number of floating point operations performed during model inference.
- torchscan.modules.module_flops(module, input, output)[source]¶
Estimate the number of floating point operations performed by the module
- Parameters:
module (torch.nn.Module) – PyTorch module
input (torch.Tensor) – input to the module
output (torch.Tensor) – output of the module
- Returns:
number of FLOPs
- Return type:
int
MACs¶
Related to the number of multiply-accumulate operations performed during model inference
- torchscan.modules.module_macs(module, input, output)[source]¶
Estimate the number of multiply-accumulation operations performed by the module
- Parameters:
module (torch.nn.Module) – PyTorch module
input (torch.Tensor) – input to the module
output (torch.Tensor) – output of the module
- Returns:
number of MACs
- Return type:
int
DMAs¶
Related to the number of direct memory accesses during model inference
- torchscan.modules.module_dmas(module, input, output)[source]¶
Estimate the number of direct memory accesses by the module. The implementation overhead is neglected
- Parameters:
module (torch.nn.Module) – PyTorch module
input (torch.Tensor) – input to the module
output (torch.Tensor) – output of the module
- Returns:
number of DMAs
- Return type:
int
Receptive field¶
Related to the effective receptive field of a layer
- torchscan.modules.module_rf(module, input, output)[source]¶
Estimate the spatial receptive field of the module
- Parameters:
module (torch.nn.Module) – PyTorch module
input (torch.Tensor) – input to the module
output (torch.Tensor) – output of the module
- Returns:
receptive field int: effective stride int: effective padding
- Return type:
int