torchinsight 0.1.1

Enhanced model analysis tool for PyTorch models

TorchInsight

TorchInsight is an enhanced PyTorch model analysis tool that provides functionality similar to torchinfo but with custom formatting and additional features.

Features

• Detailed model structure visualization
• Automatic FLOPS calculation with appropriate unit selection (K, M, G)
• Support for input dimension specification without batch dimension
• Support for long dtype specification for specific inputs
• Analysis of various model architectures (CNN, Attention, Recommendation systems, etc.)
• Colorized output for improved readability

Installation

pip install torchinsight

Quick Start

import torch
import torch.nn as nn
from torchinsight import analyze_model

# Create a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(16 * 16 * 16, 10)
        
    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(-1, 16 * 16 * 16)
        x = self.fc(x)
        return x

# Create model instance
model = SimpleModel()

# Analyze model
summary = analyze_model(
    model,
    model_name="SimpleModel",
    input_dims=(3, 32, 32),  # Input dimensions (channels, height, width)
    batch_size=64,  # Batch size
)

# Print analysis results
print(summary)

Advanced Usage

TorchInsight supports multiple input formats and data types:

# Analyze model with multiple inputs
summary = analyze_model(
    model,
    model_name="ComplexModel",
    input_dims=[(13,), (5,)],  # Two inputs with dimensions (13,) and (5,)
    long_indices=[1],  # Second input (index 1) should be torch.long
    batch_size=128,  # Batch size
)

For more examples, see the examples directory.

Documentation

For complete documentation, visit:

• Usage Guide
• API Reference

Contributing

Contributions are welcome! Feel free to submit Pull Requests or create Issues.

License

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

中文