dimviz 0.1.0

A lightweight debugging tool for tracking tensor shape transformations in PyTorch

DimViz 🔍

A lightweight debugging tool for tracking tensor shape transformations in PyTorch models.

Stop guessing tensor shapes—see them flow through your model in real-time! DimViz helps you debug shape mismatches, understand model architecture, and optimize tensor operations.

✨ Features

• 🎯 Zero Code Changes - Context manager and decorator patterns
• 📊 Rich Visualization - Beautiful terminal tables (with optional Rich library)
• 🔍 Smart Filtering - Track only what matters
• 💾 Multiple Export Formats - JSON, CSV, TXT
• 📈 Memory Tracking - See memory allocation per operation
• 🎨 Friendly Names - Human-readable operation names
• ⚡ Performance Aware - Minimal overhead with smart filtering

🚀 Installation

pip install dimviz

For enhanced visualization with colors and styling:

pip install dimviz[rich]

📖 Quick Start

Basic Usage

import torch
import torch.nn as nn
from dimviz import DimViz

# Your model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# Track shape transformations
x = torch.randn(32, 784)

with DimViz():
    output = model(x)

Output:

[DimViz] 🟢 Tracking Started...
[DimViz] 🔴 Tracking Finished. (Elapsed: 0.02s)

╭──────┬────────────┬─────────────┬──────────────╮
│ Step │ Operation  │ Input Shape │ Output Shape │
├──────┼────────────┼─────────────┼──────────────┤
│    1 │ linear     │ (32, 784)   │ (32, 256)    │
│    2 │ relu       │ (32, 256)   │ (32, 256)    │
│    3 │ linear     │ (32, 256)   │ (32, 10)     │
╰──────┴────────────┴─────────────┴──────────────╯

[DimViz] 📊 Summary:
  • Total Operations: 3
  • Unique Operations: 2
  • Logged Operations: 3

Decorator Usage

from dimviz import visualize

@visualize()
def train_step(model, batch):
    output = model(batch)
    loss = criterion(output, target)
    loss.backward()
    return loss

loss = train_step(model, batch)

🎛️ Advanced Features

Track Only Shape Changes

When you only care about shape transformations:

with DimViz(verbose=False):
    output = model(x)

This filters out operations that don't change tensor shapes (like activation functions on the same tensor).

Memory Tracking

Monitor memory allocation per operation:

with DimViz(track_memory=True):
    output = model(x)

Output includes memory columns:

╭──────┬───────────┬─────────────┬──────────────┬─────────┬──────────╮
│ Step │ Operation │ Input Shape │ Output Shape │ Mem In  │ Mem Out  │
├──────┼───────────┼─────────────┼──────────────┼─────────┼──────────┤
│    1 │ linear    │ (32, 784)   │ (32, 256)    │ 0.10MB  │ 0.03MB   │
│    2 │ relu      │ (32, 256)   │ (32, 256)    │ 0.03MB  │ 0.03MB   │
│    3 │ linear    │ (32, 256)   │ (32, 10)     │ 0.03MB  │ 0.00MB   │
╰──────┴───────────┴─────────────┴──────────────┴─────────┴──────────╯

[DimViz] 📊 Summary:
  • Peak Memory: 0.10MB
  • Total Memory Delta: -0.07MB

Filter Specific Operations

Focus on particular operation types:

with DimViz(filter_ops=['conv2d', 'matmul', 'linear']):
    output = model(x)

Export Logs

Save your debugging session for later analysis:

from dimviz import export_log

with DimViz() as viz:
    output = model(x)

# Export to various formats
export_log(viz.get_log(), 'debug_log.json')
export_log(viz.get_log(), 'debug_log.csv')
export_log(viz.get_log(), 'debug_log.txt')

Compare Model Runs

Compare shape flows between different model versions:

from dimviz.exporter import compare_logs

# First model
with DimViz() as viz1:
    output1 = model_v1(x)

# Second model
with DimViz() as viz2:
    output2 = model_v2(x)

# Compare
diff = compare_logs(viz1.get_log(), viz2.get_log(), "v1", "v2")
print(diff)

Use Cases

1. Debugging Shape Mismatches

# Find where dimensions go wrong
with DimViz():
    x = torch.randn(32, 3, 224, 224)
    x = conv1(x)  # (32, 64, 112, 112)
    x = conv2(x)  # (32, 128, 56, 56)
    x = x.view(32, -1)  # See the flattened size!
    x = fc(x)  # Does it match?

2. Understanding Complex Architectures

# Trace transformer attention mechanism
with DimViz(verbose=False):  # Only shape changes
    attention_output = transformer_layer(x)

3. Optimizing Memory Usage

# Find memory-hungry operations
with DimViz(track_memory=True, filter_ops=['conv2d', 'linear']):
    output = large_model(x)

4. Teaching & Documentation

# Generate shape flow documentation
@visualize(verbose=True)
def forward_pass(x):
    """Documented forward pass with shape tracking."""
    return model(x)

🔧 Configuration Options

DimViz()

Parameter	Type	Default	Description
verbose	bool	True	Log all operations (True) or only shape changes (False)
track_memory	bool	False	Track memory allocation per operation
filter_ops	List[str]	None	Only track specific operations
max_entries	int	None	Limit number of logged operations
show_summary	bool	True	Display summary statistics

@visualize()

Same parameters as DimViz(), used as a decorator:

@visualize(verbose=False, track_memory=True)
def my_function(x):
    return model(x)

📊 Supported Operations

DimViz translates PyTorch operations to friendly names:

• aten::mm → matmul
• aten::addmm → linear_proj
• aten::conv2d → conv2d
• aten::bmm → batch_matmul
• aten::cat → concat
• aten::sigmoid → sigmoid
• ...and many more!

Operations are automatically detected and logged with human-readable names.

🎨 Output Formats

Terminal Output

With Rich (if installed):

• Colored, styled tables
• Better readability
• Automatic terminal adaptation

Without Rich:

• ASCII tables via tabulate
• Still clear and readable
• Works everywhere

Export Formats

JSON - Structured data with metadata:

{
  "metadata": {
    "timestamp": "2024-01-29T10:30:00",
    "total_operations": 10
  },
  "log": [
    {
      "step": 1,
      "operation": "linear",
      "input_shape": "(32, 784)",
      "output_shape": "(32, 256)"
    }
  ]
}

CSV - Spreadsheet-friendly:

step,operation,input_shape,output_shape
1,linear,"(32, 784)","(32, 256)"
2,relu,"(32, 256)","(32, 256)"

TXT - Human-readable logs:

DimViz Log Export
================================================================================
Step | Operation | Input Shape | Output Shape
1    | linear    | (32, 784)   | (32, 256)
2    | relu      | (32, 256)   | (32, 256)

⚡ Performance

DimViz uses PyTorch's TorchDispatchMode for minimal overhead:

• Verbose mode: ~5-10% slowdown (logs everything)
• Non-verbose mode: ~2-5% slowdown (logs only shape changes)
• With filtering: ~1-3% slowdown (logs specific ops only)

For production code, use verbose=False or filter_ops to minimize impact.

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

Development Setup

git clone https://github.com/yourhimalaya-kaushik/dimviz.git
cd dimviz
pip install -e ".[dev]"
pytest tests/

Running Tests

pytest tests/ -v
pytest tests/ --cov=dimviz --cov-report=html

📝 License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• Built on PyTorch's TorchDispatchMode API
• Visualization powered by Rich (optional)
• Table formatting by tabulate

📧 Contact

• GitHub: @himalaya-kaushik
• Email: [e-mail]

---

Made with ❤️ for the PyTorch community

If you find DimViz helpful, please consider giving it a ⭐ on GitHub!