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: Module, inputs: Tuple[Tensor, ...], out: Tensor) → int[source]

Estimate the number of floating point operations performed by the module

Parameters:

  • module – PyTorch module

  • inputs – input to the module

  • out – output of the module

Returns:

number of FLOPs

MACs

Related to the number of multiply-accumulate operations performed during model inference

torchscan.modules.module_macs(module: Module, inp: Tensor, out: Tensor) → int[source]

Estimate the number of multiply-accumulation operations performed by the module

Parameters:

  • module (torch.nn.Module) – PyTorch module

  • inp (torch.Tensor) – input to the module

  • out (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: Module, inp: Tensor, out: Tensor) → int[source]

Estimate the number of direct memory accesses by the module. The implementation overhead is neglected.

Parameters:

  • module (torch.nn.Module) – PyTorch module

  • inp (torch.Tensor) – input to the module

  • out (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: Module, inp: Tensor, out: Tensor) → Tuple[float, float, float][source]

Estimate the spatial receptive field of the module

Parameters:

  • module (torch.nn.Module) – PyTorch module

  • inp (torch.Tensor) – input to the module

  • out (torch.Tensor) – output of the module

Returns:

receptive field effective stride effective padding