ONNX Graph Visualization

Relevant source files

The following files were used as context for generating this wiki page:

• src/components/OnnxGraphVisualizer.tsx
• src/utils/modelInfo.ts
• src/WorkbenchContext.tsx
• src/types.ts
• src/components/CenterCanvas.tsx
• src/components/Sidebar.tsx

ONNX Graph Visualization is a core feature of the ONNX-Station workbench designed to provide developers with high-fidelity insights into model architecture, operator distribution, and memory utilization. It intercepts ONNX operator streams during the compilation pipeline to visualize how different optimization levels and hardware targets affect the resulting graph topology.

The system focuses on two primary metrics: operator frequency (instance counts) and estimated VRAM memory allocation. This allows users to identify computational bottlenecks and heavy-weight layers that may impact performance on edge devices or web-based environments.

Sources: src/components/OnnxGraphVisualizer.tsx:342-358, src/utils/modelInfo.ts:5-8

System Architecture and Data Flow

The visualization system operates as a reactive consumer of the WorkbenchContext. When a model is "compiled" (either via a simulated trace or actual Optimum export), the system generates structural metadata that is then rendered using D3.js.

Data Generation Pipeline

The following diagram illustrates the flow from user configuration to the final visual representation:

The data flow starts with ExportConfig, which defines the model path, modality, and optimization level. The generateCompiledModelInfo utility processes these inputs to simulate or extract operator statistics.
Sources: src/WorkbenchContext.tsx:188-210, src/utils/modelInfo.ts:11-30

Core Components

1. Model Info Utility (modelInfo.ts)

This utility serves as the logic engine for the visualizer. It calculates high-fidelity simulated distributions based on the model's modality (text, vision, audio) and optimization parameters.

Variable	Description	Impact on Visualization
numLayers	Determined by model path (e.g., Whisper Large = 24).	Increases operator counts and memory scale.
dtypeScale	Scale factor (fp16 = 0.50, fp32 = 1.0).	Directly reduces estimated memory Mb.
optOpCountMultiplier	Scaling factor based on O1-O4 optimization.	O3/O4 reduces counts via fused kernels.

Sources: src/utils/modelInfo.ts:14-65

2. Graph Visualizer Component (OnnxGraphVisualizer.tsx)

The visualizer component uses D3.js to render SVG-based charts. It responds to window resizing and user hover interactions to provide detailed information about specific operator types.

• Bar Chart: Displays operator instance counts (frequency). Compute-intensive operators like MatMul or Conv are categorized and colored differently from shape-related operators like Reshape.
• Donut Chart: Represents estimated VRAM highlights in MB. It identifies the "Peak VRAM Memory Allocation" bottleneck, often associated with heavyweight layer weights.

Sources: src/components/OnnxGraphVisualizer.tsx:68-150, src/components/OnnxGraphVisualizer.tsx:365-385

Operator Categorization Logic

The system categorizes ONNX operators into functional groups to help developers understand the graph's computational nature.

The color coding applied in OnnxGraphVisualizer.tsx maps directly to these categories for visual consistency.
Sources: src/components/OnnxGraphVisualizer.tsx:42-50, src/utils/modelInfo.ts:80-150

Implementation Details

Operator Data Structure

The visualizer relies on the CompiledModelOperator type, which encapsulates the metrics for a specific operator type within the graph.

// src/types.ts:60-66
export interface CompiledModelOperator {
  opType: string;
  count: number;
  memoryMb: number;
  percentage: number;
  category: 'compute' | 'memory' | 'activation' | 'normalization' | 'shape';
}

Optimization Mapping

The visualization changes dynamically based on the optimize flag (O1 through O4). For instance, at O3 (Complete Attention Fusion), standard MatMul and Softmax nodes are replaced by MultiHeadAttention (Fused) nodes in the visualizer output.

// src/utils/modelInfo.ts:114-117
if (config.optimize === 'O3') {
  addOp('MultiHeadAttention (Fused)', numLayers, numLayers * 6.0 * dtypeScale * optMemMultiplier, 'compute');
  addOp('Softmax (Remainder)', Math.round(numLayers * 0.2), numLayers * 0.4 * dtypeScale, 'activation');
}

Sources: src/utils/modelInfo.ts:114-117, src/types.ts:60-66

Summary

The ONNX Graph Visualization module provides a critical bridge between abstract configuration and hardware-level execution reality. By profiling operator node frequency and peak VRAM allocation, it enables developers to verify graph topology and dynamic axes before deployment. The integration of D3.js ensures that complex structural data is accessible through interactive, responsive charts.