ONNX Go Parser

Чистый, надёжный, модульный парсер ONNX-файлов на Go без CGO.

Описание

ONNX Go Parser - это библиотека на языке Go для парсинга ONNX-моделей. Она предоставляет AST (абстрактное синтаксическое дерево) графа модели и веса в виде []float32, готовые к генерации кода. Библиотека автоматически определяет, хранятся ли веса в одном файле или во внешних файлах данных (.onnx.data), и корректно загружает их в обоих случаях. Библиотека разработана с философией "только stdlib, zero-dependency", что делает её идеальной для production-сред.

Особенности

• Чистый Go код без CGO
• Нулевые внешние зависимости
• Поддержка широкого спектра ONNX-операций (математические, активации, свертки, нормализация, условные и др.)
• Поддержка весов и тензоров
• Поддержка атрибутов узлов
• Совместимость с современными ONNX-моделями (включая MiniLM, BERT, CNN и другие)
• Протестировано на реальной модели MiniLM: 260 узлов, 507 рёбер, 191 инициализатор, 3 входа, 1 выход
• Поддержка квантования (FP32, FP16, INT8, UINT8) через вспомогательный конвертер
• Построение вычислительного графа с топологической сортировкой
• Визуализация графа в формате DOT для Graphviz
• Поддержка циклических зависимостей в архитектурах моделей
• Поддержка связей между узлами (рёбер графа)
• Инструменты командной строки для анализа моделей
• Вывод информации о модели в различных форматах (JSON, ASCII, DOT)
• Автоматическое определение и загрузка весов как из одного файла, так и из внешних файлов данных (.onnx.data)
• Поддержка моделей любого размера: от маленьких CNN до больших трансформеров

Установка

go mod init your-project
go get git.ymnuktech.ru/ymnuk/go-onnx-parser

Использование

Базовое использование

package main

import (
    "fmt"
    "log"

    onnxparser "git.ymnuktech.ru/ymnuk/go-onnx-parser"
)

func main() {
    // Загрузка ONNX-модели
    graph, err := onnxparser.LoadModel("path/to/model.onnx")
    if err != nil {
        log.Fatal("Ошибка при загрузке модели:", err)
    }

    fmt.Printf("Модель: %s\n", graph.Name)
    fmt.Printf("Количество узлов: %d\n", len(graph.Nodes))
    fmt.Printf("Количество весов: %d\n", len(graph.Initializers))

    // Обработка узлов
    for _, node := range graph.Nodes {
        fmt.Printf("Узел: %s (%s)\n", node.OpType, node.Name)
        // Обработка атрибутов и связей
    }

    // Пример: получение конкретного тензора по имени
    // Библиотека автоматически загружает веса как из одного файла, так и из внешних файлов (.onnx.data)
    tensor, exists := graph.GetTensorByName("weight_name")
    if exists {
        fmt.Printf("Тензор: %s, форма: %v, тип: %d\n", tensor.Name, tensor.Shape, tensor.DataType)
        // Обработка данных тензора (весов)
        if len(tensor.FloatData) > 0 {
            fmt.Printf("Первые 5 значений: %v\n", tensor.FloatData[:min(5, len(tensor.FloatData))])
        }
        
        // Использование метода GetFloatData для получения весов в виде []float32
        // независимо от исходного типа данных (FLOAT, INT8, UINT8, FP16 и т.д.)
        floatWeights := tensor.GetFloatData()
        fmt.Printf("Веса в виде []float32: %v\n", floatWeights[:min(5, len(floatWeights))])
    }

    // Пример: получение конкретного узла по имени
    node, exists := graph.GetNodeByName("node_name")
    if exists {
        fmt.Printf("Узел: %s, тип: %s\n", node.Name, node.OpType)
        // Обработка узла
    }

    // Пример: получение всех узлов определенного типа операции
    convNodes := graph.GetNodeByOpType("Conv")
    fmt.Printf("Найдено %d сверточных узлов\n", len(convNodes))
    // Обработка сверточных узлов
}

Структуры данных

Библиотека предоставляет следующие основные структуры:

// Graph представляет вычислительный граф модели
type Graph struct {
    Name         string
    Inputs       []ValueInfo
    Outputs      []ValueInfo
    Nodes        []Node
    Initializers map[string]Tensor // веса
    Edges        []Edge            // связи между узлами
}

// Node представляет узел в вычислительном графе
type Node struct {
    Name       string
    OpType     string
    Inputs     []string
    Outputs    []string
    Attributes map[string]interface{}
}

// Edge представляет связь между узлами в вычислительном графе
type Edge struct {
    Source      string // Имя исходного узла или инициализатора
    Destination string // Имя целевого узла
    OutputName  string // Имя выходного тензора из источника
    InputName   string // Имя входного тензора в приемник
}

// ExternalDataInfo содержит информацию о внешнем расположении данных тензора
type ExternalDataInfo struct {
    Offset int64  // Смещение в файле данных
    Length int64  // Длина данных в байтах
    File   string // Имя файла данных
}

// Tensor представляет тензор данных (веса, константы)
type Tensor struct {
    Name            string            // Имя тензора
    Shape           []int64           // Размерности тензора
    DataType        int               // Тип данных ONNX (1 для FLOAT, 2 для UINT8, и т.д.)
    FloatData       []float32         // Значения данных с плавающей запятой
    Int64Data       []int64           // Значения данных Int64
    RawData         []byte            // Необработанные байтовые данные для других типов
    ExternalData    *ExternalDataInfo // Информация о внешнем расположении данных (если используется)
}

// GetFloatData возвращает веса в виде []float32, независимо от исходного типа данных
//
// Этот метод преобразует веса из исходного типа данных (INT8, UINT8, FP16 и т.д.) 
// в формат float32, который можно использовать для вычислений.
func (t *Tensor) GetFloatData() []float32

// ValueInfo представляет информацию о входе/выходе
type ValueInfo struct {
    Name  string
    Type  string
    Shape []int64
}

Работа с вычислительным графом

Библиотека предоставляет методы для работы с вычислительным графом:

// Топологическая сортировка узлов
sortedNodes, err := graph.TopologicalSort()
if err != nil {
    // Возможна ошибка, если граф содержит циклы
    log.Printf("Граф содержит циклы: %v", err)
}

// Экспорт графа в формат DOT для визуализации
err = graph.ExportDOT("graph.dot")
if err != nil {
    log.Printf("Ошибка экспорта: %v", err)
}
// Затем можно использовать Graphviz для визуализации:
// dot -Tpng graph.dot -o graph.png

// Получение конкретного тензора по имени
tensor, exists := graph.GetTensorByName("weight_name")
if exists {
    fmt.Printf("Тензор: %s, форма: %v, тип: %d\n", tensor.Name, tensor.Shape, tensor.DataType)
    // Обработка данных тензора (весов)
}

// Получение конкретного узла по имени
node, exists := graph.GetNodeByName("node_name")
if exists {
    fmt.Printf("Узел: %s, тип: %s\n", node.Name, node.OpType)
    // Обработка узла
}

// Получение всех узлов определенного типа операции
convNodes := graph.GetNodeByOpType("Conv")
fmt.Printf("Найдено %d сверточных узлов\n", len(convNodes))
// Обработка сверточных узлов

Конвертер ONNX моделей

Проект включает в себя мощный Python-скрипт для конвертации различных форматов моделей в ONNX:

Установка зависимостей

cd converter/
pip install -r requirements.txt

Использование конвертера

python convert_to_onnx.py [ПУТЬ_К_МОДЕЛИ] -o [ПУТЬ_К_ВЫХОДНОМУ_ФАЙЛУ] [ОПЦИИ]

Опции

• INPUT: Путь к входному файлу или директории модели
• -o OUTPUT, --output OUTPUT: Путь к выходному ONNX файлу (обязательный параметр)
• --input-shape SHAPE: Форма входных данных для PyTorch моделей (например, --input-shape 1 3 224 224)
• --precision PRECISION: Точность представления весов (по умолчанию fp32):
  • fp32 - 32-битные числа с плавающей запятой (стандартный формат)
  • fp16 - 16-битные числа с плавающей запятой (половинная точность)
  • int8 - 8-битные целые числа со знаком
  • uint8 - 8-битные целые числа без знака
• --dynamic-axes DYNAMIC_AXES: Определение динамических осей (переменной длины), например:
  • 'input:batch_size,output:batch_size' - простой формат
  • 'input:{0:\'batch_size\'},output:{0:\'seq_len\'}' - расширенный формат
  • 'input.0:batch_size=4,output.0:batch_size=4' - формат с фиксированными значениями
  • 'input.0:batch_size=1,input.1:seq_len=128' - фиксированные значения для разных осей

Примеры использования

Конвертировать PyTorch модель с квантованием в FP16:

python convert_to_onnx.py model.pth -o model_fp16.onnx --input-shape 1 3 224 224 --precision fp16

Конвертировать модель с фиксированной длиной последовательности (например, для LSTM):

python convert_to_onnx.py lstm_model.pth -o lstm_fixed.onnx --input-shape 1 10 512 --dynamic-axes 'input:{0:"batch_size",1:"seq_len"},output:{0:"batch_size"}' --precision fp16

Конвертировать LSTM модель с фиксированными размерами batch и sequence:

python convert_to_onnx.py lstm_model.pth -o lstm_fixed.onnx --input-shape 4 50 256 --dynamic-axes 'input.0:batch_size=4,input.1:seq_len=50' --precision fp16

Подробное объяснение параметров

--input-shape

Параметр --input-shape определяет форму (размерность) входных данных для вашей модели. Он необходим для PyTorch моделей, чтобы скрипт мог создать "фиктивный" тензор правильной формы и передать его в модель для экспорта в ONNX.

Что это такое:

• Это список целых чисел, представляющих размеры каждого измерения входного тензора
• Каждое число соответствует размеру одного измерения

Как выбрать значения:

• Первое значение обычно размер батча (batch size) - количество образцов, обрабатываемых одновременно
• Последующие значения - размеры признаков, высота/ширина изображения, длина последовательности и т.д.

Откуда взять:

1. Из кода обучения/вывода модели: Найдите место, где вы передаете данные в модель, и посмотрите размеры тензора
2. Из документации модели: Некоторые предобученные модели имеют документацию с описанием требуемой формы входа
3. Из примеров использования: В репозиториях часто есть примеры с правильной формой входа

Примеры:

• Для изображений (1 образец, 3 канала RGB, 224x224 пикселей): --input-shape 1 3 224 224
• Для текста (1 батч, 128 токенов, 768 признаков): --input-shape 1 128 768
• Для видео (2 видео, 3 канала, 16 кадров, 224x224): --input-shape 2 3 16 224 224

Как рассчитать:

1. Определите, сколько измерений имеет ваш вход (обычно 2-5)
2. Установите размер батча (обычно 1 для инференса, может быть больше для тренировки)
3. Узнайте размеры остальных измерений из документации модели или кода

--dynamic-axes

Параметр --dynamic-axes позволяет указать, какие размеры в форме тензора могут изменяться во время инференса. Это особенно важно для моделей, работающих с переменными длинами последовательностей (например, RNN, трансформеры) или переменными размерами батчей.

Что это такое:

• Это способ указать ONNX, какие оси тензора могут иметь переменный размер
• Позволяет модели обрабатывать данные разного размера без перекомпиляции

Как выбрать значения:

• batch_size - используется для оси, которая может изменяться в зависимости от размера батча
• seq_len - используется для оси, которая может изменяться в зависимости от длины последовательности
• Любые другие пользовательские имена для других переменных размеров

Откуда взять:

1. Из архитектуры модели: Если модель может обрабатывать переменные длины последовательностей, ось последовательности должна быть динамической
2. Из требований к инференсу: Если вы планируете использовать переменные размеры батчей, ось батча должна быть динамической
3. Из документации ONNX: Официальная документация содержит примеры для различных архитектур

Примеры:

• Простой случай: 'input:batch_size,output:batch_size' - позволяет изменять размер батча
• Для трансформеров: 'input_ids:{0:"batch_size",1:"seq_len"},attention_mask:{0:"batch_size",1:"seq_len"}' - позволяет изменять как размер батча, так и длину последовательности
• С фиксированными значениями: 'input.0:batch_size=4,input.1:seq_len=128' - фиксирует размер батча в 4 и длину последовательности в 128

Как рассчитать:

1. Определите, какие оси вашего тензора могут изменяться (обычно это ось батча и/или ось последовательности)
2. Назначьте каждому индексу оси подходящее имя (batch_size, seq_len и т.д.)
3. Укажите, какие из них должны быть динамическими, а какие фиксированными

Поддерживаемые форматы

• PyTorch (.pth, .pt, .ckpt)
• TensorFlow SavedModel (директория содержащая saved_model.pb)
• Keras (.h5, .keras)
• ONNX (.onnx) - прямая передача с возможностью квантования

Особенности экспорта ONNX-моделей

При экспорте больших моделей может происходить разделение на два файла:

• model.onnx - содержит структуру модели (граф, узлы, связи)
• model.onnx.data - содержит веса модели

Это происходит автоматически, когда размер модели превышает определенный порог. Такое разделение помогает избежать проблем с ограничениями файловой системы на размер отдельных файлов.

Для создания единой модели (в одном файле) можно использовать следующий подход:

import onnx

# Загрузить модель
model = onnx.load("model.onnx")

# Убедиться, что все тензоры встроены в модель
for init in model.graph.initializer:
    if init.HasField('data_location') and init.data_location == onnx.TensorProto.EXTERNAL:
        # Если тензор использует external data, преобразуем его
        raw_data = onnx.numpy_helper.to_array(init)
        init.ClearField('external_data')
        init.raw_data = raw_data.tobytes()

# Сохранить модель в одном файле
onnx.save(model, "model-unified.onnx", save_as_external_data=False)

Если модель экспортируется в два файла, для корректной работы с ней необходимо убедиться, что оба файла находятся в одной директории.

Как определить форму входа для различных типов моделей

CNN (Сверточные нейронные сети)

Для CNN, обрабатывающих изображения:

• Обычно форма: [batch_size, channels, height, width] или [batch_size, height, width, channels] в зависимости от фреймворка
• channels: 1 для градационных изображений, 3 для RGB, 4 для RGBA
• height, width: размеры изображения (например, 224x224, 299x299, 384x384)

Пример: --input-shape 1 3 224 224 для одного RGB изображения 224x224

RNN/LSTM/GRU (Рекуррентные нейронные сети)

Для рекуррентных сетей:

• Обычно форма: [batch_size, sequence_length, features] или [sequence_length, batch_size, features]
• sequence_length: длина последовательности (например, количество слов в предложении)
• features: размерность признаков (например, размер эмбеддинга)

Пример: --input-shape 1 50 256 для одного батча, 50 токенов, 256 признаков на токен

Трансформеры

Для трансформеров:

• Обычно форма: [batch_size, sequence_length, hidden_size]
• hidden_size: размер скрытого состояния (например, 768 для BERT-base, 1024 для BERT-large)
• Также могут быть дополнительные входы, такие как attention_mask

Пример: --input-shape 1 128 768 для одного батча, 128 токенов, 768 признаков

MLP (Многослойные перцептроны)

Для многослойных перцептронов:

• Обычно форма: [batch_size, input_features]
• input_features: количество входных признаков

Пример: --input-shape 32 784 для батча из 32 образцов с 784 признаками каждый

Структура проекта

onnx-go-parser/
├── parser.go          // Основной парсер
├── graph.go           // AST: Graph, Node, Tensor
├── ops/               // Реестр поддерживаемых операций
│   └── supported.go
├── proto/             // Сгенерированные Go-структуры из onnx.proto
├── converter/         // Скрипт конвертации в ONNX
│   ├── convert_to_onnx.py
│   ├── requirements.txt
│   └── README.md
├── model/             // Директория для ONNX моделей
├── docs/              // Документация
└── README.md

Требования

• Go 1.25+
• Python 3.7+ (для конвертера)

Пример использования

В директории examples/ находится пример использования библиотеки:

cd examples/
go run example_usage.go

Пример демонстрирует:

• Загрузку ONNX-модели
• Вывод информации о графе
• Попытку топологической сортировки
• Экспорт графа в DOT-формат для визуализации

Инструменты командной строки

Проект включает утилиту командной строки для анализа ONNX-моделей:

# Базовая информация о модели
go run cmd/onnx-inspect/main.go model.onnx

# Вывод в формате JSON
go run cmd/onnx-inspect/main.go --json model.onnx

# Генерация DOT-файла для Graphviz
go run cmd/onnx-inspect/main.go --dot model.onnx

# Простая ASCII-визуализация
go run cmd/onnx-inspect/main.go --ascii model.onnx

# Проверка поддерживаемых операций (whitelist)
go run cmd/onnx-inspect/main.go --check-ops model.onnx

Утилита предоставляет:

• Список входов/выходов модели
• Количество узлов и уникальных операций
• Список уникальных операций с количеством узлов
• Размеры весов модели
• Возможность экспорта в различные форматы для анализа
• Проверку соответствия операций в модели белому списку поддерживаемых операций

Проверка корректности парсера

Для проверки корректности нашей реализации мы предоставляем возможность сравнения результатов с эталонной библиотекой ONNX:

# Сравнение результатов парсинга с эталонной библиотекой ONNX
python3 scripts/compare_inspectors.py model.onnx

Этот скрипт сравнивает результаты, полученные нашим Go-парсером, с результатами эталонной библиотеки ONNX, выводя таблицу сравнения количества узлов для каждой операции:

Operation                 Reference (ONNX)   Go Parser    Match
------------------------------------------------------------
Add                       123                123          ✓    
Cast                      2                  2            ✓    
Div                       12                 12           ✓    
...

Проверка поддерживаемых операций

Для проверки, все ли операции из белого списка поддерживаются в указанной модели, используйте флаг --check-ops:

# Проверка, какие операции из белого списка присутствуют в модели
go run cmd/onnx-inspect/main.go --check-ops model.onnx

Эта команда выведет:

• Общее количество поддерживаемых операций в белом списке
• Количество операций, присутствующих в модели
• Список поддерживаемых операций, которые НЕ используются в модели
• Список операций, которые присутствуют в модели, но НЕ поддерживаются (не находятся в белом списке)

Пример вывода:

Model: main_graph
Supported Operations Check:
  Total supported operations in whitelist: 69
  Operations present in model: 15
  Supported operations present in model: 15
  Supported operations NOT present in model: 54
  Unexpected operations (not in whitelist): 0

  Supported operations NOT present in model:
    - Abs
    - And
    - ArgMax
    ...

Это позволяет быстро определить, поддерживает ли ваша модель все необходимые операции, и какие операции из белого списка не используются в конкретной модели.

Проверка корректности

Мы провели всестороннюю проверку корректности нашей реализации путем сравнения результатов с эталонной библиотекой ONNX на нескольких типах моделей:

Сравнение с эталонной библиотекой ONNX

Для различных моделей:

Модель	Метрика	ONNX библиотека (эталон)	Go-парсер
bert-ner.onnx	Nodes	469	469
	Initializers	211	211
	Inputs	2	2
	Outputs	1	1
	Unique Operations	15	15
	Total parameters	107,823,522	107,823,522
minilm-l6.onnx	Nodes	260	260
	Initializers	191	191
	Inputs	3	3
	Outputs	1	1
	Unique Operations	21	21
	Total parameters	22,566,055	22,566,055
tiny-cnn.onnx	Nodes	10	10
	Initializers	9	9
	Inputs	1	1
	Outputs	1	1
	Unique Operations	5	5
	Total parameters	268,652	268,652
lstm.onnx	Nodes	28	28
	Initializers	16	16
	Inputs	1	1
	Outputs	1	1
	Unique Operations	8	8
	Total parameters	52,874	52,874

Для проверки используйте предоставленные инструменты:

# Сравнение результатов парсинга с эталонной библиотекой ONNX
python3 scripts/compare_inspectors.py model/bert-ner.onnx

# Проверка с Go-парсером
go run cmd/onnx-inspect/main.go model/bert-ner.onnx

# Проверка поддерживаемых операций
go run cmd/onnx-inspect/main.go --check-ops model/bert-ner.onnx

# Автоматический тест всех моделей (Stage 6)
bash scripts/test_stage6.sh

Подробное сравнение операций

Вот таблица сравнения количества узлов для каждой операции в модели bert-ner.onnx:

Operation                 Reference (ONNX)   Go Parser    Match
------------------------------------------------------------
Add                       123                123          ✓
Cast                      2                  2            ✓
Div                       12                 12           ✓
Erf                       12                 12           ✓
Expand                    1                  1            ✓
Gather                    2                  2            ✓
LayerNormalization        25                 25           ✓
MatMul                    97                 97           ✓
Mul                       48                 48           ✓
Reshape                   72                 72           ✓
Softmax                   12                 12           ✓
Sub                       1                  1            ✓
Transpose                 60                 60           ✓
Unsqueeze                 1                  1            ✓
Where                     1                  1            ✓
------------------------------------------------------------
Total operations: 15
Reference (ONNX) found: 469 total nodes
Go parser found: 469 total nodes
Mismatches: No

Как видно из таблицы, все операции и их количество узлов полностью совпадают между эталонной библиотекой ONNX и нашим Go-парсером, что подтверждает корректность нашей реализации.

Тестирование различных типов моделей

Мы протестировали наш парсер на следующих типах моделей:

• MiniLM-L6: Модель трансформера для задач NLP
• BERT-NER: Модель для Named Entity Recognition
• Tiny CNN: Простая сверточная нейросеть для классификации изображений
• LSTM: Рекуррентная нейросеть для обработки временных рядов

Все модели были успешно загружены, проанализированы и результаты проверены с эталонной библиотекой ONNX.

Документация

• Структура вычислительного графа - подробное описание узлов и ребер графа
• ROADMAP - план разработки проекта

Лицензия

MIT