TensorKiko
_______ _ _______ _ ______
|__ __| | |/ /_ _| |/ / __ \
| | ___ _ __ ___ ___ _ __| ' / | | | ' / | | |
| |/ _ \ '_ \/ __|/ _ \| '__| < | | | <| | | |
| | __/ | | \__ \ (_) | | | . \ _| |_| . \ |__| |
|_|\___|_| |_|___/\___/|_| |_|\_\_____|_|\_\____/
TensorKiko is a powerful and intuitive tool for visualizing and analyzing machine learning model structures. It supports various model formats and provides detailed insights into model architecture, parameters, and tensor statistics.
Features
- Multi-format Support: Load and process models in .safetensors, .pt, .pth, .pb, and .h5 formats.
- Interactive Visualization: Generate HTML-based visualizations of model structures with a tree-based layout.
- Detailed Analysis:
- Model parameters, memory usage, and estimated FLOPs.
- Tensor statistics (mean, std dev, min/max values, zero count) with histograms.
- SVG representation of tensor shapes.
- Anomaly Detection: Automatically detect and highlight potential issues in model tensors.
- Search Functionality: Easily navigate large models.
- Custom Layer Filtering: Include or exclude specific layers using regex patterns.
- Precision Information: Display the data type of model parameters.
- Web-based Interface: User-friendly interface with collapsible sections.
UI Examples
Model Overview
This image shows the main interface, displaying:
- Model name and overall details
- Breakdown of layer types
- Beginning of the model's hierarchical structure
Layer Details
This image demonstrates the detailed view of a specific layer, including:
- Hierarchical model structure
- Layer information (parameters, shape, statistics)
- Histogram of weight distribution
Installation and Usage
Requirements
Installation
pip install tensorkiko
Basic Usage
tensorkiko path/to/your/model.safetensors
For multiple input files:
tensorkiko path/to/model1.pt path/to/model2.safetensors path/to/model3.h5
Command-line Options
--debug
: Enable debug mode--no-tree
: Disable tree visualization--port PORT
: Specify HTTP server port (default: 8000)--output-dir DIR
: Set output directory--include-layers REGEX
: Include only specific layers--exclude-layers REGEX
: Exclude specific layers
Example:
tensorkiko path/to/model.pt --debug --port 8080 --include-layers "conv|linear"
Web Interface Guide
- Run TensorKiko on your model(s).
- A web browser will open with the visualization.
- Explore the collapsible header for model information and layer type statistics.
- Click tree nodes to view detailed layer information.
- Use the search bar to find specific layers or parameters.
- Examine tensor statistics, histograms, and shape visualizations.
- Check for highlighted anomalies in the interface.
Contributing
We welcome contributions! Please submit a Pull Request on our GitHub repository.
License
This project is licensed under the MIT License. See the LICENSE file for details.
TensorKiko is developed by the open-source community and takara.ai staff. The project is sponsored by takara.ai.