CircuitGen_AI

Аннотация

CircuitGen_AI это исследовательский ML-проект для предсказания характеристик цифровых схем по их графовому представлению. В текущем состоянии проект поддерживает пайплайн parser -> dataloader -> graph model -> trainer и ориентирован на задачи регрессии, в первую очередь на предсказание delay и area.

Проект работает с двумя типами входных данных:

• GraphML-графы с метаданными в JSON.
• Verilog-описания схем, из которых граф строится во время preprocessing.

Технически проект опирается на PyTorch, torch-geometric, networkx и набор собственных абстракций для парсеров, загрузчиков данных, моделей, тренеров и критериев.

Цель

Цель проекта - предоставить воспроизводимый каркас для экспериментов с графовыми нейросетями на данных о цифровых схемах.

На практике проект нужен для решения следующих задач:

• предсказание целевых характеристик схем без запуска полного внешнего EDA-процесса;
• сравнение разных GNN-подходов на едином pipeline;
• быстрое подключение новых форматов данных, новых моделей и новых trainer-классов;
• исследование preprocessing-стратегий для графов схем.

Какие проблемы решает проект

• Схемы хранятся в разных представлениях, а модели ожидают единый формат данных. Проект приводит входные данные к torch_geometric.data.Data.
• Эксперименты с GNN часто распадаются на несвязанные скрипты. Здесь есть единая точка сборки пайплайна через MainEngine.
• Датасеты могут содержать типовые ошибки. Для GraphML-ветки предусмотрены автоматические фиксы и проверки.
• Для разных моделей нужны разные trainer-реализации и функции потерь. Проект выделяет их в отдельные расширяемые слои.

В чем суть проекта

Суть проекта в том, что граф схемы рассматривается как основной объект обучения. Дальше:

1. Парсер читает исходный формат данных.
2. DataLoader превращает данные в объекты PyG Data.
3. Trainer собирает модель, criterion и optimizer.
4. Обучение запускается через единый orchestration-слой.
5. Артефакты запуска сохраняются в Run_result.

Сейчас основная точка входа - main.py, которая создаёт Scripts/Main_Engine.py и запускает конфиг Start_Configs/Test.json.

Архитектура

Верхний уровень

• main.py - текущий entrypoint проекта.
• Scripts/Main_Engine.py - orchestration pipeline: чтение JSON-конфига, создание parser, dataloader и trainer, запуск Train или Test.
• Start_Configs - примеры конфигов запуска.
• Trainer_and_Models - модели, trainer-классы, callback и специализированные training loop.
• preprocessing - парсеры, dataloader-слой, preprocessing графов и фиксы ошибок.
• Criterions - функции потерь и фабрика критериев.
• docker и .gitlab-ci.yml - инфраструктура сборки и запуска.

Поддерживаемые сценарии данных

GraphML

• Парсер: GraphMLParser
• DataLoader: GraphMLDataLoader
• Цели: delay, area через abcStatsBalance[target]
• Преобразование: GraphML -> networkx -> one-hot node features -> PyG Data

Verilog

• Парсер: VerilogParser
• DataLoader: VerilogDataLoader
• Преобразование: Verilog -> networkx -> Node2Vec embeddings -> PyG Data
• Поддержка сейчас заточена под конкретный формат датасета и naming convention

Поддерживаемые trainer-классы

• TransformerTrainer - регрессия на базе TransformerGNN
• GraphConvTrainer - регрессия на базе GCNModel
• HOGA_Trainer - отдельный trainer для HOGA-подобной модели с high-order preprocessing

Дополнительно в репозитории есть заготовки под MLP_Trainer, ContrastiveTrainer и Ray Tune, но они не встроены в основной пользовательский flow так же явно, как три сценария выше.

Быстрый старт

1. Требования

• Python 3.10 или 3.11
• Желательно Linux-окружение или Docker
• Для GPU-сценариев: CUDA 11.8

pyproject.toml уже ограничивает версию Python диапазоном >=3.10,<3.12.

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

Минимальный набор:

pip install -r requirements.txt

Для полного GNN-стека также нужны torch, torch-geometric, torch-scatter, torch-sparse. В Docker они ставятся отдельным скриптом docker/install-heavy-deps.sh.

Если вы настраиваете окружение вручную вне Docker, ориентируйтесь на:

• pyproject.toml
• requirements.txt
• requirements-tune.txt
• docker/install-heavy-deps.sh

3. Подготовка конфига

Основной пример конфига находится в Start_Configs/Test.json.

Ключевые секции:

• parser
• data_loader
• trainer
• random_seed
• test_or_train
• experiment_name

Пример важных полей:

{
  "test_or_train": "Train",
  "parser": {
    "parser_class": "GraphMLParser",
    "target": "delay"
  },
  "data_loader": {
    "data_loader_class": "GraphMLDataLoader",
    "data_path": "../CircuitGen_AI/datasets/dataset_with_errors",
    "batch_size": 2,
    "encoder": "One-hot-encoder",
    "shuffle": true,
    "val_size": 0.2,
    "n_workers": 20,
    "use_scaler": false
  },
  "trainer": {
    "trainer_class": "TransformerTrainer",
    "learning_rate": 0.001,
    "device": "cuda",
    "criterion": "MAPE_loss"
  }
}

4. Запуск

Текущий запуск:

python main.py

Важно: main.py сейчас жёстко использует Start_Configs/Test.json. Если нужен другой конфиг, есть два варианта:

• временно изменить путь внутри main.py;
• создать MainEngine программно с нужным путём к JSON.

5. Результаты

Во время обучения callback создаёт директорию запуска в ../CircuitGen_AI/Run_result и сохраняет туда:

• best_model.pth
• config.json
• log.txt
• TensorBoard-логи

Как пользоваться проектом

Базовый сценарий работы такой:

1. Подготовить датасет в поддерживаемом формате.
2. Выбрать parser и dataloader в JSON-конфиге.
3. Выбрать trainer и criterion.
4. Запустить обучение через python main.py.
5. Анализировать метрики и чекпоинты в Run_result.

Для тестового режима:

1. В конфиге указать "test_or_train": "Test".
2. Передать путь к директории run через experiment_name.
3. Убедиться, что в ней есть best_model.pth.

Раздел для инженеров ML - NN

Что здесь уже есть

• готовый pipeline для graph regression;
• базовые GNN-модели на TransformerConv и GCNConv;
• отдельная HOGA-подобная ветка с high-order preprocessing;
• CriterionFactory для выбора loss из конфига;
• логирование, ранняя остановка и сохранение лучшего checkpoint через TrainerCallback.

Где смотреть основные точки

• Trainer_and_Models/Models.py - архитектуры моделей
• Trainer_and_Models/Model_trainer.py - базовый training loop
• Trainer_and_Models/HOGA_Trainer.py - отдельный trainer для HOGA
• Criterions/custom_criterion.py - loss-функции
• preprocessing/preprocessing.py - graph preprocessing для расширенных признаков

Как добавлять новую модель

1. Добавить модель в Trainer_and_Models/Models.py или отдельный модуль.
2. Если модель требует стандартный интерфейс, реализовать forward(x, edge_index, batch).
3. Добавить trainer-класс, наследующий AbstractModelTrainer.
4. Подключить этот trainer в Scripts/Main_Engine.py, чтобы его можно было выбрать по имени в JSON.
5. Добавить или переиспользовать criterion.

Как добавлять новый loss

1. Реализовать класс в Criterions/custom_criterion.py.
2. Зарегистрировать его в Criterions/CriterionFactory.py.
3. Использовать имя класса в секции trainer.criterion JSON-конфига.

Ограничения текущей ML-части

• HOGA_Trainer использует собственную training/evaluation-логику и не полностью следует общей схеме criterion.
• VerilogDataLoader сейчас жёстко привязан к конкретному датасету и шаблону имён файлов.
• Часть экспериментальных компонентов в репозитории ещё не интегрирована в основной запуск.

Раздел для разработчиков

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

Структура расширения проекта

Новый parser

Реализуйте наследника AbstractDataParser и зарегистрируйте его в preprocessing/Parser/ParserFactory.py.

Новый dataloader

Реализуйте наследника AbstractDataLoader или AbstractGraphDataLoader и зарегистрируйте его в preprocessing/Dataloader/DataloaderFactory.py.

Новый encoder

Добавьте реализацию в encoders.py и убедитесь, что data_loader.encoder может выбрать её по имени.

Новый trainer

Реализуйте trainer и подключите его к MainEngine, чтобы выбор происходил через trainer_class в JSON-конфиге.

Docker

Базовый образ описан в Dockerfile.

Особенности:

• базовый runtime - nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04;
• тяжёлые зависимости PyTorch/PyG ставятся через отдельный скрипт;
• поддержан режим deferred install через переменную DEFER_HEAVY_DEPS.

Стандартный контейнерный запуск по умолчанию вызывает:

python main.py

CI

Текущий GitLab CI описан в .gitlab-ci.yml.

Сейчас он делает следующее:

• запускает заглушку tests;
• собирает Docker image через docker buildx;
• пушит образ в приватный registry.

Известные ограничения

• README теперь синхронизирован с кодом, но сам проект всё ещё содержит экспериментальные и частично недоведённые ветки.
• main.py не принимает путь к конфигу через CLI.
• Не все зависимости задекларированы в одном источнике истины: часть разбросана между pyproject.toml, requirements.txt и docker-скриптами.
• Тесты в CI пока не настроены.
• Часть комментариев и старой документации в репозитории всё ещё содержит проблемы кодировки.

Авторы

Список авторов хранится в docs/AUTHORS.md.