torchscan

Crawler

torchscan.crawl_module(module, input_shape, dtype=None, max_depth=None)

Retrieves module information for an expected input tensor shape

Example::

>>> 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 (torch.nn.Module) – module to inspect

• input_shape (tuple<int>) – expected input shapes

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

• max_depth (int, optional) – maximum depth of layer information

Returns:

layer and overhead information

Return type:

dict

torchscan.summary(module, input_shape, wrap_mode='mid', max_depth=None, receptive_field=False)

Print module summary for an expected input tensor shape

Example::

>>> 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 (torch.nn.Module) – module to inspect

• input_shape (tuple<int>) – expected input shapes

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

• max_depth (int, optional) – maximum depth of layer information

• receptive_field (bool, optional) – whether receptive field estimation should be performed