get_module[source]

get_module(o, i)

Recursively get the module from list of indices

Most of the times, we want to extract nested module in architecture; while some modules support indexing (Sequential), some don't. This function enables you to access nested modules in a numpy-like indexing. When coupled with arch_summary, we can effortlessly explore the pytorch models.

arch_summary[source]

arch_summary(arch, idx=None, verbose=False)

Short architecture summary, used for holistic understanding and deciding parameter groups

from torchvision.models import shufflenet_v2_x1_0, resnet50

Examples

arch_summary can handle named as well as simple Sequential modules with different behavior. Indices could be used to quickly navigate the architecture, module names are also introduced in summary if available. xresnet50 is the architecture defined in fastai2 and is extending Sequential

arch_summary(xresnet50)

[0 ] ConvLayer        : 3   layers
[1 ] ConvLayer        : 3   layers
[2 ] ConvLayer        : 3   layers
[3 ] MaxPool2d        : 1   layers
[4 ] Sequential       : 3   layers
[5 ] Sequential       : 4   layers
[6 ] Sequential       : 6   layers
[7 ] Sequential       : 3   layers
[8 ] AdaptiveAvgPool2d: 1   layers
[9 ] Flatten          : 1   layers
[10] Dropout          : 1   layers
[11] Linear           : 1   layers

get_module could be used to extract a specific module nested deep within hierarchy.

get_module(xresnet50,[4,0])

ResBlock(
  (convpath): Sequential(
    (0): ConvLayer(
      (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (1): ConvLayer(
      (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU()
    )
    (2): ConvLayer(
      (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (idpath): Sequential(
    (0): ConvLayer(
      (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (act): ReLU(inplace=True)
)

And if you really want to go deeper, you may set verbose=True and arch_summary will go 2 depth down. For the simplicity, I'm keeping it to the depth of 2, since you can always have a detailed summary using fastai2's patched summary method on module

arch_summary(xresnet50,verbose=True)

[0 ] ConvLayer        : 3   layers
      Conv2d
      BatchNorm2d
      ReLU
[1 ] ConvLayer        : 3   layers
      Conv2d
      BatchNorm2d
      ReLU
[2 ] ConvLayer        : 3   layers
      Conv2d
      BatchNorm2d
      ReLU
[3 ] MaxPool2d        : 1   layers
[4 ] Sequential       : 3   layers
      ResBlock
      ResBlock
      ResBlock
[5 ] Sequential       : 4   layers
      ResBlock
      ResBlock
      ResBlock
      ResBlock
[6 ] Sequential       : 6   layers
      ResBlock
      ResBlock
      ResBlock
      ResBlock
      ResBlock
      ResBlock
[7 ] Sequential       : 3   layers
      ResBlock
      ResBlock
      ResBlock
[8 ] AdaptiveAvgPool2d: 1   layers
[9 ] Flatten          : 1   layers
[10] Dropout          : 1   layers
[11] Linear           : 1   layers

Now let's quickly review a named_module architecture and how arch_summary handles that information

arch_summary(shufflenet_v2_x1_0)

[0 ] (conv1)    Sequential       : 3   layers
[1 ] (maxpool)  MaxPool2d        : 1   layers
[2 ] (stage2)   Sequential       : 4   layers
[3 ] (stage3)   Sequential       : 8   layers
[4 ] (stage4)   Sequential       : 4   layers
[5 ] (conv5)    Sequential       : 3   layers
[6 ] (fc)       Linear           : 1   layers

Refer arch_explore for more examples