torchscan

Crawler

torchscan.crawl_module(module: Module, input_shape: List[Tuple[int, ...]] | Tuple[int, ...], dtype: dtype | Iterable[dtype] | None = None) → Dict[str, Any]

Retrieves module information for an expected input tensor shape

>>> import torch.nn as nn
>>> from torchscan import summary
>>> mod = nn.Conv2d(3, 8, 3)
>>> module_info = crawl_module(mod, (3, 224, 224))

Parameters:

• module – module to inspect

• input_shape – expected input shapes

• dtype – data type of each input argument to the module

Returns:

layer and overhead information

torchscan.summary(module: Module, input_shape: Tuple[int, ...], wrap_mode: str = 'mid', max_depth: int | None = None, receptive_field: bool = False, effective_rf_stats: bool = False) → None

Print module summary for an expected input tensor shape

>>> import torch.nn as nn
>>> from torchscan import summary
>>> mod = nn.Conv2d(3, 8, 3)
>>> summary(mod, (3, 224, 224), receptive_field=True)

Parameters:

• module – module to inspect

• input_shape – expected input shapes (don’t include batch size)

• wrap_mode – if a value is too long, where the wrapping should be performed

• max_depth – maximum depth of layer information

• receptive_field – whether receptive field estimation should be performed

• effective_rf_stats – if receptive_field is True, displays effective stride and padding