GPT-Driven Neural Network Generator

_{short alias lmurg}

📖 Overview

This Python-based NNGPT project leverages large language models (LLMs) to automate the creation of neural network architectures, streamlining the design process for machine learning practitioners. It leverages various neural networks from the LEMUR Dataset to fine-tune LLMs and provide insights into potential architectures during the creation of new neural network models.

Create and Activate a Virtual Environment (recommended)

For Linux/Mac:

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install --upgrade pip

For Windows:

python3 -m venv .venv
.venv\Scripts\activate
python3 -m pip install --upgrade pip

It is assumed that CUDA 13.0 is installed; otherwise, consider replacing 'cu130' with the appropriate version. Most LLM usage scenarios require GPUs with at least 24 GB of memory.

Environment for NNGPT Developers

Pip package manager

Create a virtual environment, activate it, and run the following command to install all the project dependencies:

python -m pip install --upgrade pip
pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu130
MAX_JOBS=4 NVCC_THREADS=1 pip install -r req-no-isolation.txt --no-build-isolation --no-cache -v --extra-index-url https://download.pytorch.org/whl/cu130

Removing the --no-cache parameter allows previously built packages to be retrieved from the cache, but it can lead to problems if the PyTorch version has changed.

If there are installation problems, install the dependencies from the 'requirements.txt' file one by one.

Update of NN Dataset

To get the latest code and statistics, install the most recent version of the LEMUR Dataset from GitHub:

rm -rf db
pip uninstall -y nn-dataset
pip install --no-cache-dir git+https://github.com/ABrain-One/nn-dataset --extra-index-url https://download.pytorch.org/whl/cu130

Installing the stable version:

rm -rf db
pip uninstall -y nn-dataset 
pip install nn-dataset --extra-index-url https://download.pytorch.org/whl/cu130

Adding functionality to export data to Excel files and generate plots for analyzing neural network performance:

pip install nn-stat --extra-index-url https://download.pytorch.org/whl/cu130

and export/generate:

python -m ab.stat.export

Installation of NNGPT with pip

   pip install nn-gpt --extra-index-url https://download.pytorch.org/whl/cu130
   pip install nn-gpt[flash] --no-build-isolation --extra-index-url https://download.pytorch.org/whl/cu130

Use

• ab.gpt.NNAlter* – Generates modified neural network models.
  Use the -e argument to set the number of epochs for the initial CV model generation.

• ab.gpt.NNEval – Evaluates the models generated in the previous step.

• ab.gpt.TuneNNGen* – Performs fine-tuning and evaluation of an LLM. For evaluation purposes, the LLM generates neural network models, which are then trained to assess improvements in the LLM’s performance on this task. The -s flag allows skipping model generation for the specified number of epochs.

• ab.gpt.AccPredictor – Fine-tunes and evaluates a Qwen3-8B accuracy predictor from LEMUR training runs. Given early-epoch accuracies and neural network code, it predicts final best_accuracy and best_epoch.

  Running the script runs four steps in order:

  1. Preprocessing — loads training runs from the nn_dataset API, filters runs with ≥50 epochs, and writes ab/gpt/data/llm_finetuning_data.jsonl
  2. Data preparation for training — converts preprocessed runs into ChatML train/val/test splits (ab/gpt/data/train_llm_dataset.jsonl, val_llm_dataset.jsonl, test_llm_dataset.jsonl)
  3. Model training — QLoRA fine-tunes Qwen3-8B with validation and early stopping; saves the checkpoint to ab/gpt/model2/
  4. Model testing — runs inference on the test split and writes ab/gpt/data/test_predictions.csv and test_metrics.log

  python -m ab.gpt.AccPredictor
  

  Individual steps can also be imported:

  from ab.gpt.AccPredictor import data_preprocessing, prepare_llm_datasets, train_model, test_model, predict_best_accuracy
  from ab.gpt.AccPredictor import DEFAULT_TRAIN_PATH, DEFAULT_VAL_PATH, DEFAULT_OUTPUT_DIR, DEFAULT_TEST_PATH
  
  data_preprocessing()
  prepare_llm_datasets()
  train_model(DEFAULT_TRAIN_PATH, DEFAULT_VAL_PATH)
  test_model(model_path=DEFAULT_OUTPUT_DIR, data_path=DEFAULT_TEST_PATH)
  
  # Inference with the published Hugging Face checkpoint
  best_acc, best_epoch = predict_best_accuracy(task, dataset, metric, nn_code, epoch_1_acc, epoch_2_acc)
  

  Requires a GPU with ≥24 GB VRAM, unsloth, and the LEMUR/nn-dataset package installed.

• ab.gpt.TuneNNGen_delta.py – Delta-based fine-tuning entry point (see arXiv:2605.04903). The LLM generates compact unified diffs (deltas) to refine baseline architectures instead of full code. Uses paper-aligned hyperparameters (lr=1e-5, temperature=0.35, top-k=50, LoRA with lm_head). Calls TuneNNGen.main() with delta defaults — no upstream behavior is changed.

  python -m ab.gpt.TuneNNGen_delta
  python -m ab.gpt.TuneNNGen_delta --llm_conf qwen2.5_coder_7b_instruct.json
  

LangGraph Multi-Agent Workflow

NNGPT supports an optional LangGraph-based multi-agent orchestration mode. The agent system integrates directly inside tune() — no separate entry point, no duplicated logic.

Design Principle

All pipeline logic remains in ab/gpt/util/Tune.py as the single source of truth. Agent nodes are thin wrappers only — they read from state and call the existing functions. No logic is reimplemented inside any agent file.

Agent Flow

The professor-specified flow is: Finetuner → Generator → Evaluator → Predictor

• manager — controls routing, checks epoch stop condition, decides next node
• generator — calls nn_gen() / trans_gen(); skips if epoch < skip_epoch; skips evaluator if no code generated
• evaluator — calls _evaluate_epoch(); stores accuracy and all predictor inputs in state
• finetuner — calls _finetune_epoch(); increments epoch counter, returns to manager
• predictor — optional; activates after epoch 1 and epoch 2 accuracies are both available

Any future improvement to nn_gen(), trans_gen(), _evaluate_epoch(), or _finetune_epoch() automatically applies to both classic and agent modes.

Crash Recovery

Agent mode uses LangGraph MemorySaver checkpointing. If the pipeline crashes mid-epoch (e.g. GPU OOM), re-running with the same nn_name_prefix resumes from the last completed node — no restart from epoch 0.

Usage

The agent mode is enabled by default.

To use the accuracy predictor agent:

python -m ab.gpt.TuneNNGen --use_predictor

Agent Files

File	Purpose
ab/gpt/agents/run_agent.py	Builds and runs the LangGraph StateGraph
ab/gpt/agents/manager.py	Routing logic and epoch stop condition
ab/gpt/agents/predictor.py	Optional accuracy prediction node
ab/gpt/agents/state.py	Shared AgentState TypedDict — field names match LEMUR DB columns
ab/gpt/util/Tune.py	Single source of truth: nn_gen, trans_gen, _evaluate_epoch, _finetune_epoch, generate_step, evaluate_step, finetune_step
ab/gpt/AccPredictor.py	Accuracy predictor: data prep, fine-tuning, and evaluation

Fine-tuned LLMs

🐳 Docker

All versions of this project are compatible with AI Linux and can be seamlessly executed within the AI Linux Docker container.

Installing the latest version of the project from GitHub

docker run --rm -u $(id -u):ab -v $(pwd):/a/mm abrainone/ai-linux:llm bash -c "[ -d nn-gpt ] && git -C nn-gpt pull || git -c advice.detachedHead=false clone --depth 1 https://github.com/ABrain-One/nn-gpt"

Running script

docker run --rm -u $(id -u):ab --shm-size=16G -v $(pwd)/nn-gpt:/a/mm abrainone/ai-linux:llm bash -c "python -m ab.gpt.TuneNNGen_8B"

If recently added dependencies are missing in the AI Linux, you can create a container from the Docker image abrainone/ai-linux:llm, install the missing packages (preferably using pip install <package name>), and then create a new image from the container using docker commit <container name> <new image name>. You can use this new image locally or push it to the registry for deployment on the computer cluster.

Citation

The original version of this project was created at the Computer Vision Laboratory of the University of Würzburg by the authors mentioned below. If you find this project to be useful for your research, please consider citing our articles for NNGPT, architecture design, hyperparameter tuning and delta-based NAS with LLMs:

@InProceedings{ABrain.NNGPT,
	title = {{NNGPT}: Rethinking {AutoML} with Large Language Models},
	author = {Kochnev, Roman and Khalid, Waleed and Uzun, Tolgay Atinc and Zhang, Xi and Dhameliya, Yashkumar Sanjaybhai and Qin, Furui and Vysyaraju, Chandini and Duvvuri, Raghuvir and Goyal, Avi and Ignatov, Dmitry and Timofte, Radu},
	booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW)},		
	pages = {5664--5674},	
	year={2026}
}

@InProceedings{ABrain.Architect,
	title={From Memorization to Creativity: {LLM} as a Designer of Novel Neural Architectures},
	author={Khalid, Waleed and Ignatov, Dmitry and Timofte, Radu},
	booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW)},	
	pages = {3252--3261},	
	year={2026}
}

@InProceedings{ABrain.HPGPT,
	title={Optuna vs Code Llama: Are {LLMs} a New Paradigm for Hyperparameter Tuning?},
	author={Kochnev, Roman and Goodarzi, Arash Torabi and Bentyn, Zofia Antonina and Ignatov, Dmitry and Timofte, Radu},
	booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops (ICCVW)},
	url={https://openaccess.thecvf.com/content/ICCV2025W/AIM/papers/Kochnev_Optuna_vs_Code_Llama_Are_LLMs_a_New_Paradigm_for_ICCVW_2025_paper.pdf},
	pages = {5664--5674},
	year={2025},
	doi={10.1109/ICCVW69036.2025.00598}
}

@Article{ABrain.DeltaNAS,
	title={Delta-Based Neural Architecture Search: {LLM} Fine-Tuning via Code Diffs},
	author={Adhikari, Santosh Premi and Timofte, Radu and Ignatov, Dmitry},
	journal={arXiv preprint arXiv:2605.04903},
	year={2026}
}

Licenses

This project is distributed under the following licensing terms:

• models with pretrained weights under the legacy DeepSeek LLM V2 license
• all neural network models and their weights not covered by the above licenses, as well as all other files and assets in this project, are subject to the MIT license

The idea and leadership of Dr. Ignatov