TorchVista

https://sachinhosmani.github.io/torchvista/
Admin
TorchVista is an interactive tool to visualize PyTorch model forward passes directly in notebooks...
Tokens:21,172
Snippets:128
Trust Score:9.7
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install TorchVista

Source: https://sachinhosmani.github.io/torchvista/index.html

Install the torchvista package using pip. This command should be run in your terminal or notebook environment.

```bash
pip install torchvista
```

--------------------------------

### Trace SqueezeNet Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/SqueezeNet.html

Use this snippet to trace a SqueezeNet model with a specified example input and tracing depth. Ensure torchvision and torch are installed.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.squeezenet1_0(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=4, collapse_modules_after_depth=2)
```

--------------------------------

### Trace MiniUnetPP Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MiniUnetPP.html

Defines a MiniUnetPP model and traces it using TorchVista. Ensure the model and an example input are provided.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class UNetPP(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv00 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv10 = nn.Conv2d(16, 32, 3, padding=1)
        self.conv01 = nn.Conv2d(16 + 32, 16, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.final = nn.Conv2d(16, 1, 1)

    def forward(self, x):
        x00 = torch.relu(self.conv00(x))
        x10 = torch.relu(self.conv10(self.pool(x00)))
        x10_up = self.up(x10)
        x01 = torch.relu(self.conv01(torch.cat([x00, x10_up], dim=1)))
        return self.final(x01)

model = UNetPP()
example_input = torch.randn(1, 1, 64, 64)

trace_model(model, example_input)
```

--------------------------------

### Trace PyTorch Model with Example Input

Source: https://sachinhosmani.github.io/torchvista/index.html

Use the trace_model function from torchvista to visualize your defined PyTorch model. Provide the model instance and an example input tensor to generate the visualization.

```python
example_input = torch.randn(1, 10, 32)
trace_model(model, example_input)
```

--------------------------------

### Trace Model with TorchVista for Compression Visualization

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest6.html

Instantiates the CompressionTest6 model and uses torchvista.trace_model to visualize its structure with compression options enabled. Requires the model and an example input.

```python
model = CompressionTest6()
example_input = torch.randn(2, 16, 256)  # (batch, seq_len, d_model)

trace_model(model, example_input, collapse_modules_after_depth=0, show_compressed_view=True)
```

--------------------------------

### Trace MobilenetV3Small Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MobilenetV3Small.html

Imports necessary libraries and traces a MobilenetV3Small model with a specified input and tracing depth. Ensure torchvision is installed.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.mobilenet_v3_small(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=5)

```

--------------------------------

### Define and Trace DeepCNN Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/DeepCNN.html

Defines a DeepCNN model and then traces it using torchvista. Ensure torch and torchvista are installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.block(x)

class DeepCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.stage1 = ConvBlock(3, 32)
        self.stage2 = ConvBlock(32, 64)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(64, 10)

    def forward(self, x):
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.classifier(x)

model = DeepCNN()
example_input = torch.randn(1, 3, 64, 64)

trace_model(model, example_input)
```

--------------------------------

### Trace Model with Parameters using TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/ParameterNodes.html

This snippet shows how to define a PyTorch model with parameters and submodules, and then trace it using TorchVista for visualization. Ensure Torch and TorchVista are installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class SubModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 64)
        self.scale = nn.Parameter(torch.ones(64))

    def forward(self, x):
        return self.linear(x) * self.scale

class ModelWithParameters(nn.Module):
    def __init__(self):
        super().__init__()
        self.bias = nn.Parameter(torch.zeros(64))
        self.sub = SubModule()

    def forward(self, x):
        x = x + self.bias
        x = x + torch.ones(64)
        x = torch.relu(x)
        x = self.sub(x)
        return x

model = ModelWithParameters()
example_input = torch.randn(1, 64)

trace_model(model, example_input)
```

--------------------------------

### Define and Trace SimpleMultiPathModule with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MultiPathWithError.html

Defines a PyTorch module with multiple paths and parameters, then traces it using TorchVista. This example may produce an error during execution, which TorchVista can help visualize.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class SimpleMultiPathModule(nn.Module):
    def __init__(self, in_features=10, out_features=5):
        super(SimpleMultiPathModule, self).__init__()
        
        self.linear1 = nn.Linear(in_features, out_features)
        
        self.path_weights = nn.ParameterDict({
            'path1_weight': nn.Parameter(torch.tensor(0.5)),
        })
        
        self.global_biases = nn.ParameterList([
            nn.Parameter(torch.randn(1)),
            nn.Parameter(torch.randn(1))
        ])
    
    def forward(self, x, x2, x3):
        
        output1 = self.linear1(x)
        x += x2

        x *= F.relu(x3)
        
        results = torch.zeros(x.size(0), self.linear1.out_features, device=x.device)
                
        bias = 5 + self.global_biases[0] + self.global_biases[1] + output1[:,:]
        results = results + bias + output1[-5]
        
        return results

model = SimpleMultiPathModule()

x = torch.randn(3, 10)
y = torch.randn(3, 10)
z = torch.randn(3, 10)
example_input = (x, y, z)


trace_model(model, example_input)



```

--------------------------------

### Trace EfficientNetB0 Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/EfficientNetB0.html

Use this snippet to trace an EfficientNetB0 model with a specified input and tracing depth. Ensure torch and torchvision are installed.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.efficientnet_b0(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=7)

```

--------------------------------

### Trace ConvNeXt Tiny Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/ConvextTiny.html

Use this snippet to trace a ConvNeXt-tiny model with a specified input and tracing depth. Ensure PyTorch and TorchVista are installed.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.convnext_tiny(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=5)

```

--------------------------------

### Trace AlexNet Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/AlexNet.html

Use this snippet to trace an AlexNet model with TorchVista. Ensure you have torch, torchvista, and torchvision installed. The model is initialized without pre-trained weights.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.alexnet(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input)
```

--------------------------------

### Trace MLP with Skip Connections using TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MLPSkipConnections.html

Defines a DeepMLP model with skip connections and traces it using TorchVista. Ensure torchvista is installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class DeepMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(64, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 64)
        self.out = nn.Linear(64, 10)

    def forward(self, x):
        x1 = torch.relu(self.fc1(x))
        x2 = torch.relu(self.fc2(x1))
        x3 = torch.relu(self.fc3(x2 + x1))
        return self.out(x3)

model = DeepMLP()
example_input = torch.randn(1, 64)

trace_model(model, example_input)
```

--------------------------------

### Trace MobileViT with Collapsed Depth

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MobileVitCollapseDepth2.html

Use this snippet to trace a PyTorch model, specifically a MobileViT model from the timm library, and collapse modules after a certain depth for simplified visualization. Ensure torch, torchvision, and torchvista are installed.

```python
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)
```

--------------------------------

### Trace XLNetBaseCased Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/XLNetBaseCased.html

Use this snippet to trace the XLNetBaseCased model from Hugging Face with TorchVista. Ensure you have torch, transformers, and torchvista installed.

```python
import torch
from transformers import XLNetModel
from torchvista import trace_model

model = XLNetModel.from_pretrained("xlnet-base-cased")
example_input = torch.randint(0, 32000, (1, 10))

trace_model(model, example_input)
```

--------------------------------

### Define and Trace SimpleResNet with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/SimpleResnet.html

This snippet defines a basic Residual Block and a SimpleResNet model using PyTorch. It then instantiates the model and an example input, and finally uses `torchvista.trace_model` to generate a visualization of the model's structure and execution flow.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Conv2d(3, 16, 3, padding=1)
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = torch.randn(1, 3, 32, 32)

trace_model(model, example_input)
```

--------------------------------

### Trace and Visualize PyTorch Model

Source: https://sachinhosmani.github.io/torchvista/generated/tutorials/basic_usage.html

Use this snippet to visualize the forward pass of a simple PyTorch model. Ensure torch and torchvista are installed. The visualization allows panning, zooming, and inspecting module details.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 5)
        self.linear2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear1(x) + self.linear2(x)

model = LinearModel()
example_input = torch.randn(2, 10)

# Visualize the forward pass
trace_model(model, example_input)
```

--------------------------------

### Trace a Linear Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/LinearModel.html

Use `trace_model` to visualize a PyTorch `nn.Module`. Ensure the model and an example input are defined before calling the function.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = LinearModel()
example_input = torch.randn(2, 10)

trace_model(model, example_input)

```

--------------------------------

### Define and Trace TinyModelWithNestedSubmodules

Source: https://sachinhosmani.github.io/torchvista/generated/demos/TinyModelWithNestedSubmodules.html

Define a PyTorch model with nested submodules and then trace it using `torchvista.trace_model`. This is useful for visualizing complex model architectures. Ensure `torch` and `torchvista` are installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class InnerSubBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.inner_linear_with_long_name = nn.Linear(dim, dim)
        self.inner_activation_function_layer = nn.ReLU()

    def forward(self, x):
        return self.inner_activation_function_layer(
            self.inner_linear_with_long_name(x)
        )

class MidLevelBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.midlevel_linear_transformation_step = nn.Linear(dim, dim)
        self.deeply_nested_inner_block = InnerSubBlock(dim)

    def forward(self, x):
        x = self.midlevel_linear_transformation_step(x)
        return self.deeply_nested_inner_block(x)

class TinyModelWithNestedSubmodules(nn.Module):
    def __init__(self, dim=32):
        super().__init__()
        self.initial_projection_layer_with_lengthy_name = nn.Linear(dim, dim)
        self.midlevel_feature_processing_unit = MidLevelBlock(dim)
        self.final_output_mapping_layer = nn.Linear(dim, dim)

    def forward(self, x):
        x = self.initial_projection_layer_with_lengthy_name(x)
        x = self.midlevel_feature_processing_unit(x)
        return self.final_output_mapping_layer(x)

# Example
model = TinyModelWithNestedSubmodules()

trace_model(model, torch.randn(2, 32), show_module_attr_names=True, collapse_modules_after_depth=0, forced_module_tracing_depth=5)
```

--------------------------------

### Trace ResNet18 Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/ResnetSmall.html

Use this snippet to trace a ResNet18 model with a specified input and control tracing depth. Ensure torch, torchvision, and torchvista are installed.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.resnet18(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=3, collapse_modules_after_depth=2)
```

--------------------------------

### Trace MobileViT Model with Collapsed Depth 1

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MobileVitCollapseDepth1.html

Use this snippet to trace a pre-trained MobileViT model using TorchVista. The `collapse_modules_after_depth=1` argument simplifies the visualization by merging modules after the first layer. Ensure you have `torch`, `torchvision`, `timm`, and `torchvista` installed.

```python
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=1)
```

--------------------------------

### Define and Instantiate ComplexModel

Source: https://sachinhosmani.github.io/torchvista/generated/demos/ArtificiallyMadeComplexModel.html

Defines a complex PyTorch model with parallel classifier paths and auxiliary biases. Instantiate the model and prepare example inputs for tracing.

```python
self.aux_biases = nn.ParameterList([
            nn.Parameter(torch.randn(1)),
            nn.Parameter(torch.randn(1))
        ])
    
def forward(self, x, y):
        # Input shape: [batch_size, 3, height, width]
        x = self.features(x)
        x = self.adaptive_pool(x)
        x = torch.flatten(x, 1)
        
        # Process through parallel classifier paths
        main_features = self.classifier_paths['main'](x)
        aux_features = self.classifier_paths['auxiliary'](x)
        
        # Concatenate features from different paths
        combined = torch.cat([main_features + y[0, 0, 0, 0], aux_features], dim=1)
        
        # Add auxiliary biases for demonstration
        bias = self.aux_biases[0] + self.aux_biases[1]
        
        # Final classification
        return self.final_classifier(combined) + bias + 34

model = ComplexModel()
x = torch.randn(2, 3, 32, 32)
y = torch.randn(2, 3, 32, 32)

example_input = (x, y)

trace_model(model, example_input)
```

--------------------------------

### Define PyTorch Model for Visualization

Source: https://sachinhosmani.github.io/torchvista/index.html

Define your PyTorch model, including any custom layers like SelfAttention, and then instantiate the main model class. This setup is necessary before tracing the model with TorchVista.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        attn = torch.softmax((q @ k.transpose(-2, -1)) / self.scale, dim=-1)
        return attn @ v

class AttentionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Linear(32, 64)
        self.attn = SelfAttention(64)
        self.classifier = nn.Linear(64, 5)

    def forward(self, x):
        x = self.embed(x)
        x = self.attn(x)
        return self.classifier(x.mean(dim=1))

model = AttentionClassifier()
```

--------------------------------

### Demonstrate Tensor Transpose Properties with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/TensorTransposeProperties.html

This code defines a PyTorch module that showcases the .T, .mT, and .H tensor transpose properties. It then uses TorchVista to trace the model with example inputs, including 2D, 3D batched, and complex tensors.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class TransposePropertiesModel(nn.Module):
    """
    Demonstrates the three tensor transpose properties: .T, .mT, and .H
    - .T: Simple transpose (swaps all dimensions for 2D, reverses all dims for nD)
    - .mT: Matrix transpose (transposes last two dimensions, works on batches)
    - .H: Hermitian transpose (conjugate transpose for complex tensors)
    """
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(8, 8)
        self.linear2 = nn.Linear(8, 8)
        self.linear3 = nn.Linear(8, 8, dtype=torch.cfloat)

    def forward(self, x):
        # .T - simple 2D transpose
        y1 = self.linear1(x)
        out1 = y1 @ y1.T

        # .mT - batch matrix transpose (transposes last 2 dims)
        y2 = self.linear2(x)
        out2 = y2 @ y2.mT

        # .H - Hermitian (conjugate) transpose for complex tensors
        y3 = self.linear3(x)
        out3 = y3 @ y3.H

        return out1, out2.mean(dim=0), out3

model = TransposePropertiesModel()

example_input = (
    torch.randn(4, 8),              # 2D input for .T
    torch.randn(2, 4, 8),           # 3D batch input for .mT
    torch.randn(4, 8, dtype=torch.cfloat),  # complex input for .H
)

trace_model(model, example_input)

```

--------------------------------

### Trace Model with Repeated Outputs and Nested Structures

Source: https://sachinhosmani.github.io/torchvista/generated/demos/RepeatedOutputTensors.html

This example defines custom classes and a PyTorch module with a forward pass that returns a dictionary with nested custom objects and repeated tensor operations. Use this to visualize complex model outputs.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class Container0:
    def __init__(self, x, y):
        self.x = x
        self.y = y

class Container:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.container0 = Container0(x, y)

class CM(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(4, 4)

    def forward(self, x):
        y = self.linear1(x)
        return {0: self.linear1(self.linear1(x) + 3), 1: {0: self.linear1(x) + torch.ones(2,4)}, 2: Container(x - 2, y - 2)}

  
class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(4, 4)
        self.cm = CM()

    def forward(self, x):
        y = self.cm(x)[0]
        return Container(x, y)

model = CustomModel()

example_input = torch.randn(2, 4)

trace_model(model, example_input)
```

--------------------------------

### Trace Nested Modules with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/tutorials/exploring_modules.html

This code defines `ResidualBlock` and `SimpleResNet` modules, then uses `torchvista.trace_model` to visualize the model structure, including nested `nn.Sequential` and custom blocks. Ensure PyTorch and TorchVista are installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        # This Sequential can be expanded in the visualization
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Conv2d(3, 16, 3, padding=1)
        # These nested modules can be expanded too
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = torch.randn(1, 3, 32, 32)

trace_model(model, example_input)
```

--------------------------------

### Define and Trace MiniInception Model

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CNNWithBranching.html

Defines a MiniInception model with multiple convolutional branches and then traces it using TorchVista. Ensure PyTorch and TorchVista are installed. This snippet is useful for visualizing complex model architectures.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class MiniInception(nn.Module):
    def __init__(self):
        super().__init__()
        self.branch1 = nn.Conv2d(3, 8, 1)
        self.branch3 = nn.Conv2d(3, 8, 3, padding=1)
        self.branch5 = nn.Conv2d(3, 8, 5, padding=2)
        self.pool = nn.Conv2d(3, 8, 1)
        self.final = nn.Conv2d(32, 10, 1)

    def forward(self, x):
        b1 = torch.relu(self.branch1(x))
        b3 = torch.relu(self.branch3(x))
        b5 = torch.relu(self.branch5(x))
        bp = torch.relu(self.pool(torch.max_pool2d(x, 3, stride=1, padding=1)))
        return self.final(torch.cat([b1, b3, b5, bp], dim=1))

model = MiniInception()
example_input = torch.randn(1, 3, 32, 32)

trace_model(model, example_input)
```

--------------------------------

### Trace Faster R-CNN MobileNetV3 Large 320 FPN Model

Source: https://sachinhosmani.github.io/torchvista/generated/demos/FasterrcnnMobilenetV3Large320Fpn.html

This snippet loads a pre-trained Faster R-CNN model, prepares an example input, and then uses `torchvista.trace_model` to trace the model. Adjust `forced_module_tracing_depth` and `collapse_modules_after_depth` for detailed or collapsed tracing.

```python
import torch
from torchvista import trace_model
from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_320_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

model = fasterrcnn_mobilenet_v3_large_320_fpn(pretrained=True)
example_input = [torch.rand(3, 320, 320)]

model.eval()

trace_model(model, example_input, show_non_gradient_nodes=True, forced_module_tracing_depth=5, collapse_modules_after_depth=0)
```

--------------------------------

### Trace CompressionTest7 Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest7.html

Traces the CompressionTest7 model with a sample input using TorchVista. Options `collapse_modules_after_depth=0` and `show_compressed_view=True` are used to control the visualization detail.

```python
model = CompressionTest7()
example_input = torch.randn(2, 16, 128)  # (batch, seq_len, dim)

trace_model(model, example_input, collapse_modules_after_depth=0, show_compressed_view=True)
```

--------------------------------

### Trace and Visualize Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest11.html

Traces the CompressionTest11 model with a sample input and visualizes the compressed view. This function helps in understanding module repetition and model structure.

```python
model = CompressionTest11()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True, collapse_modules_after_depth=2)
```

--------------------------------

### Trace and Visualize Compressed Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest1.html

Use this snippet to trace a PyTorch model and display its compressed view. Ensure torch and torch.nn are imported, and the model is defined. The `show_compressed_view=True` argument enables the compressed visualization.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest1(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(64, 64),
            nn.Linear(64, 64),
            nn.Linear(64, 64),
            nn.Linear(64, 64),
        )

    def forward(self, x):
        return self.layers(x)

model = CompressionTest1()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)

```

--------------------------------

### Define a Complex Model Architecture

Source: https://sachinhosmani.github.io/torchvista/generated/demos/ArtificiallyMadeComplexModel.html

A complex model integrating ResidualBlocks, ParallelPathways, AttentionModules, and DynamicLayerSelector. This serves as an example of composing various advanced modules for a sophisticated task.

```python
class ComplexModel(nn.Module):
    def __init__(self):
        super(ComplexModel, self).__init__()
        
        # Feature extraction layers
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            ResidualBlock(32, 64, stride=2),
            ParallelPathways(64, 128),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        # Adaptive pooling to get fixed size regardless of input
        self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
        
        # Classifier with multiple paths
        self.classifier_paths = nn.ModuleDict({
            'main': nn.Sequential(
                nn.Linear(128, 256),
                nn.ReLU(),
                AttentionModule(256)
            ),
            'auxiliary': DynamicLayerSelector(128, 256)
        })
        
        # Final classification layer
        self.final_classifier = nn.Linear(256 * 2, 10)
        
        # Parameter lists for demonstration
```

--------------------------------

### Visualize Model Compression with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest3.html

Use `trace_model` with `show_compressed_view=True` to visualize the compressed view of a PyTorch model. This is useful for understanding how TorchVista optimizes and compresses the model's structure.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest3(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
            ),
            nn.Sequential(
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
            ),
            nn.Sequential(
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
            ),
            nn.Sequential(
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
                nn.Linear(64, 64),
            ),
        ])

    def forward(self, x):
        for seq in self.layers:
            x = seq(x)
        return x

model = CompressionTest3()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)
```

--------------------------------

### Define and Trace SelfAttentionClassifier with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/SelfAttentionClassifier.html

This snippet defines a custom PyTorch model with a SelfAttention mechanism and then uses TorchVista's `trace_model` function to generate a visualization. Ensure you have `torch` and `torchvista` installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        attn = torch.softmax((q @ k.transpose(-2, -1)) / self.scale, dim=-1)
        return attn @ v

class AttentionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Linear(32, 64)
        self.attn = SelfAttention(64)
        self.classifier = nn.Linear(64, 5)

    def forward(self, x):
        x = self.embed(x)
        x = self.attn(x)
        return self.classifier(x.mean(dim=1))

model = AttentionClassifier()
example_input = torch.randn(2, 10, 32)

trace_model(model, example_input)
```

--------------------------------

### Define and Trace LeNet5 Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/LeNet5.html

Defines a LeNet5 convolutional neural network and then uses torchvista.trace_model to generate a visualization. Ensure torch and torch.nn are imported.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class LeNet5(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 6, 5), nn.Tanh(), nn.AvgPool2d(2),
            nn.Conv2d(6, 16, 5), nn.Tanh(), nn.AvgPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(16*4*4, 120), nn.Tanh(),
            nn.Linear(120, 84), nn.Tanh(),
            nn.Linear(84, num_classes)
        )
    def forward(self, x):
        return self.classifier(self.features(x))

model = LeNet5()
example_input = torch.randn(2, 1, 28, 28)
trace_model(model, example_input)
```

--------------------------------

### Trace PyTorch Model with Compressed View

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest4.html

Use `trace_model` to visualize a PyTorch model. Set `show_compressed_view=True` to display a compressed graph representation. Ensure necessary imports are present.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest4(nn.Module):
    def __init__(self):
        super().__init__()
        block = nn.Sequential(*[nn.Linear(64, 64) for _ in range(4)])
        self.layers = nn.ModuleList([block] * 4)

    def forward(self, x):
        for seq in self.layers:
            x = seq(x)
        return x

model = CompressionTest4()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)

```

--------------------------------

### Trace Vision Transformer Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/VisionTransformer.html

Use this snippet to trace a Vision Transformer model from torchvision. Ensure the model and example input are correctly defined. The `forced_module_tracing_depth` parameter controls the depth of module tracing.

```python
import torch
from torchvista import trace_model
from torchvision import models

model = models.vit_b_16(weights=None)
example_input = torch.randn(1, 3, 224, 224)
trace_model(model, example_input, forced_module_tracing_depth=5)
```

--------------------------------

### Trace Model with Compressed View

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest5.html

Use `trace_model` to visualize a PyTorch model. Set `show_compressed_view=True` to display a compressed representation of the model graph, which can be useful for understanding complex architectures.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest5(nn.Module):
    def __init__(self):
        super().__init__()
        block = nn.Sequential(*[nn.Linear(64, 64) for _ in range(10)])
        self.layers = nn.ModuleList([block] * 10)

    def forward(self, x):
        for seq in self.layers:
            x = seq(x)
        return x

model = CompressionTest5()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)
```

--------------------------------

### Trace Model with Controlled Expansion Depth

Source: https://sachinhosmani.github.io/torchvista/generated/tutorials/collapse_modules_after_depth.html

Use `collapse_modules_after_depth` to set the initial expansion limit for modules when tracing a model. Modules beyond this depth will be collapsed by default.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1)
        )
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = torch.randn(1, 3, 32, 32)

# Expand modules up to depth 3 when first displayed
trace_model(
    model,
    example_input,
    forced_module_tracing_depth=3,
    ##############################
    collapse_modules_after_depth=3  # <-- modules beyond this depth start collapsed
    ##############################
)
```

--------------------------------

### Define and Trace CompressionTest8 Model with TorchVista

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest8.html

Defines a PyTorch nn.Module with sequential linear layers of non-matching dimensions. Traces the model using torchvista to visualize its structure, expecting no compression due to incompatible layer sizes.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest8(nn.Module):
    """
    4 linear layers with non-matching dimensions - no compression should be detected.
    Each layer has different input/output dimensions, so they cannot be grouped.
    """
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(64, 128),
            nn.Linear(128, 256),
            nn.Linear(256, 512),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        return self.layers(x)

model = CompressionTest8()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)
```

--------------------------------

### Trace Model Hiding Non-Gradient Nodes

Source: https://sachinhosmani.github.io/torchvista/generated/tutorials/show_non_gradient_nodes.html

Use this snippet to trace a PyTorch model and hide non-gradient nodes, such as constant tensors and scalars, from the visualization. This focuses the graph on learnable parameters. Ensure `torchvista` and `torch` are installed.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1)
        )
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, inputs):
        x = inputs['image']
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = {'image': torch.randn(1, 3, 32, 32)}

# Hide constant tensors and scalars that don't require gradients
trace_model(
    model,
    example_input,
    forced_module_tracing_depth=3,
    collapse_modules_after_depth=3,
    ##############################
    show_non_gradient_nodes=False  # <-- hides non-gradient nodes
    ##############################
)
```

--------------------------------

### Define TransformerBlock and CompressionTest6 Models

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest6.html

Defines the building blocks for a Transformer model and a custom compression test model using PyTorch. Imports are required.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class TransformerBlock(nn.Module):
    def __init__(self, d_model=256, nhead=8, dim_feedforward=512):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        x = self.norm1(x + self.attn(x, x, x)[0])
        x = self.norm2(x + self.ffn(x))
        return x

class CompressionTest6(nn.Module):
    def __init__(self, num_layers=6, d_model=256, nhead=8):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, nhead) for _ in range(num_layers)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
```

--------------------------------

### Trace MobileViT Model with Module Collapsing

Source: https://sachinhosmani.github.io/torchvista/generated/demos/MobileVitCollapsed.html

Use `trace_model` to trace a MobileViT model from the timm library. Set `collapse_modules_after_depth=0` to collapse all modules after the initial layer, simplifying the graph visualization.

```python
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=0)
```

--------------------------------

### Define and Trace CompressionTest10 Model

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest10.html

Defines a PyTorch model with parallel sequential and ModuleList branches. It then traces the model using `torchvista.trace_model` to generate a visualized graph, with `show_compressed_view=True` to display a compressed representation.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class CompressionTest10(nn.Module):
    """
    Parallel branches with repeating linear layers (2 in each chain).
    Input goes to both a Sequential and a ModuleList with matching dims,
    then outputs are combined.
    """
    def __init__(self):
        super().__init__()
        self.branch_sequential = nn.Sequential(
            nn.Linear(64, 64),
            nn.Linear(64, 64),
        )
        self.branch_modulelist = nn.ModuleList([
            nn.Linear(64, 64),
            nn.Linear(64, 64),
        ])

    def forward(self, x):
        out1 = self.branch_sequential(x)
        out2 = x
        for layer in self.branch_modulelist:
            out2 = layer(out2)
        return out1 + out2

model = CompressionTest10()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)

```

--------------------------------

### Trace Model with Compressed View

Source: https://sachinhosmani.github.io/torchvista/generated/demos/CompressionTest13.html

Use `trace_model` to visualize a PyTorch model, including its parameters and structure, with a compressed view enabled. This is useful for analyzing complex models with repeated components.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class BlockWithParameter(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 64)
        self.scale = nn.Parameter(torch.ones(64))

    def forward(self, x):
        return self.linear(x) * self.scale

class CompressionTest13(nn.Module):
    """Compression test with nn.Parameter nodes inside repeated blocks."""
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            BlockWithParameter(),
            BlockWithParameter(),
            BlockWithParameter(),
            BlockWithParameter(),
        )

    def forward(self, x):
        return self.layers(x)

model = CompressionTest13()
example_input = torch.randn(2, 64)

trace_model(model, example_input, show_compressed_view=True)

```
Tokens:
Raw