TorchVista

Sample code (Not editable)

import torch
from torchvision.models import squeezenet1_1
from torchvista import trace_model

import timm

model = timm.create_model('mobilevit_s', pretrained=True)
example_input = torch.randn(1, 3, 256, 256)

trace_model(model, example_input, collapse_modules_after_depth=2, show_non_gradient_nodes=False)
Error Output (if any)
Visualized Interactive Graph
ConvNormActCollapse module-BatchNormAct2dCollapse module-SequentialCollapse module-SequentialCollapse module-SequentialCollapse module-SequentialCollapse module-SequentialCollapse module-SequentialCollapse module-ConvNormActCollapse module-BatchNormAct2dCollapse module-ClassifierHeadCollapse module-SelectAdaptivePool2dCollapse module-(1, 32, 128, 128)(1, 64, 64, 64)(1, 64, 64, 64)(1, 64, 64, 64)(1, 96, 32, 32)(1, 96, 32, 32)(1, 128, 16, 16)(1, 128, 16, 16)(1, 160, 8, 8)(1, 160, 8, 8)(1, 3, 256, 256)(1, 16, 128, 128)(1, 16, 128, 128)(1, 16, 128, 128)(1, 16, 128, 128)(1, 640, 8, 8)(1, 640, 8, 8)(1, 640, 8, 8)(1, 640, 8, 8)(1, 640, 1, 1)(1, 640)(1, 640)(1, 1000)(1, 1000)batch_normTensor OpShow infoiIdentityModuleShow infoiSiLUModuleShow infoiConv2dModuleShow infoiBottleneckBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiMobileVitBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiMobileVitBlockModuleExpand module+Show infoiBottleneckBlockModuleExpand module+Show infoiMobileVitBlockModuleExpand module+Show infoibatch_normTensor OpShow infoiIdentityModuleShow infoiSiLUModuleShow infoiConv2dModuleShow infoiAdaptiveAvgPool2dModuleShow infoiFlattenModuleShow infoiDropoutModuleShow infoiLinearModuleShow infoiIdentityModuleShow infoiinput_0output_0