Полина Люмен (OCR/DigitGPT)

Маленькая GPT-подобная модель (4.8M параметров), которая читает полоску из 1–5 цифр подряд и выписывает их в строку. Реализация — decoder-only трансформер поверх патчей картинки.

Важно про архитектуру репозитория. Это не модель для transformers: нет PreTrainedModel, нет AutoModelForCausalLM, нет tokenizers. Код написан на чистом PyTorch, веса лежат в safetensors, токенизатор — обычный класс DigitTokenizer. Для загрузки нужен код из папки scripts/ этого репозитория.

Model Details

Архитектура. DigitGPT — decoder-only трансформер с причинной маской внимания. Входная картинка режется на патчи 4×4, каждый линейно проецируется в d_model; затем к патчам приписываются токены цифр, и модель предсказывает следующую цифру авторегрессивно.

Токенизатор. Собственный, словарь из 12 токенов: 0–9, <pad> = 10, <eos> = 11. Никакого внешнего словаря.

Выход. Строка цифр без разделителей. Длина определяется самой моделью через <eos>, а не задаётся заранее.

Intended use

• Распознавание 1–5 цифр на плотной горизонтальной полоске из MNIST-подобных цифр, шириной ровно 28 px на цифру, выровненных по левому краю.
• Учебный пример GPT-подобной архитектуры, работающей с картинкой как с последовательностью: патчи + токены в одном контексте, причинная маска.
• Позиция baseline: оценка, сколько даёт маленький decoder-only при обучении на CPU.

Out of scope

Модель не рассчитана на: рукописный ввод, фотографии реальных документов, разное количество цифр вне диапазона 1–5, разделители и другие символы, цветные изображения, шум/размытие/перекос, цифры, не вошедшие в MNIST.

Training data

Датасет собирается на лету из MNIST (torchvision.datasets.MNIST), а не заранее в файл:

1. выбирается случайное число цифр k из [1, 5];
2. k цифр берётся случайно из MNIST и склеивается встык на полотно 28×140 — по 28 px на цифру, без зазоров и полей;
3. train_samples / val_samples задают количество полосок в каждой эпохе и в каждой валидации, а не размер фиксированного датасета.

Каждая эпоха поэтому показывает модель новые изображения, а не повторяет одни и те же. Модель никогда не видела реальных фотографий.

Обучение MNIST распространяется на CC BY-SA 3.0 (Yann LeCun, Corinna Cortes, Christopher J.C. Burges). Атрибуция обязательна при использовании и весов, и обучающих данных.

Training procedure

Параметры взяты из training_config.yaml, сохранённого в бандле, — это авторитетная запись запуска:

Кривая обучения — монотонная, переобучения нет:

Evaluation results

85% — это teacher-forced token accuracy на валидации, то есть доля правильно названных токенов, когда модель «посмотрела» на истинные предыдущие токены. Это мягкая метрика. Для настоящей задачи «прочитать цифры» важнее другая — точное совпадение всей строки.

Замер независимой проверкой на 200 полосках из MNIST test split (генерация с нуля, greedy=True, модель истинные токены не видит):

Разница между 0.850 и 0.828 — шум выборки, а не расхождение: при n=500 стандартная ошибка при p=0.85 равна ±1.6 пп, при n=200 — ±2.5 пп.

Как читать эти числа: модель всегда выдаёт правильное число цифр, но одну-две из пяти читает неверно. Разрыв 0.828 ↔ 0.480 — это цена одной ошибки на позицию при строгом сравнении строк. Для использования «распознать все цифры» это существенно: цифра с ошибкой делает весь результат неверным.

How to use

Зависимости: torch, safetensors, numpy, Pillow (+ PyYAML, torchvision — только для обучения):

pip install -r scripts/requirements.txt

Команда запуска — из корня репозитория:

python scripts/infer.py path/to/image.png -m .

Флаг -m . обязателен: веса и config.json лежат в корне репозитория, а по умолчанию код ищет бандл в ./model.

Опции:

python scripts/infer.py img.png -m . --sampling    # сэмплирование вместо greedy
python scripts/infer.py img.png -m . --show-input  # сохранить препроцесс 28x140
python scripts/infer.py img.png -m . --device cuda

Из Python:

import sys; sys.path.insert(0, "scripts")
from checkpoint import load_bundle
from infer import preprocess_image

model, tokenizer, manifest = load_bundle(".")        # strict=True
x = preprocess_image("img.png", 28, 140)             # сама приведёт к 28x140

ids = model.generate(x, max_len=6, greedy=True)
print(tokenizer.decode(ids[0].tolist()))             # например: '93113'

strict=True — при несовпадении весов и config.json будет исключение, а не молчаливо случайная модель.

Обучение с нуля (внимание: train.py читает config.yaml из текущего каталога, поэтому запускать из папки scripts):

cd scripts
python train.py
python train_more.py 10 --resume -m ..   # продолжить из корня

Preprocessing: что модель ожидает

infer.preprocess_image пропорционально вписывает картинку в полотно 28×140 (28 px на цифру, до 5 цифр), выравнивая по левому краю и центруя по вертикали. Увеличение картинки (4x, 10x) и уменьшение работают.

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

• широкие белые поля вокруг цифр;
• инверсия (чёрные цифры на белом фоне);
• сильный шум, размытие, перекос.

Попытка лечить это автокадрированием по чернилам была проверена и отклонена: на бенчмарке из 80 случаев точное совпадение упало с 21/80 до 6/80, потому что цифра MNIST занимает ~20 px из 28, и подгонка по высоте растягивала её до 28 px. Реализовано только LANCZOS при даунскейле — это убирает алиасинг без сдвига масштаба.

Files

README.md                    эта карточка
model.safetensors            веса (float32)
config.json                  архитектура + спецификация картинки (источник истины)
tokenizer.json               словарь: 0-9, <pad>=10, <eos>=11
training_config.yaml         полный конфиг запуска, из которого получен бандл
metrics.json                 метрики по эпохам
scripts/                     весь исходный код

В scripts/:

Limitations

• Слабое качество. 48% точных строк на 5 цифр — многовато для практики. Причина — 10000 полосок в эпоху; это маленький бюджет, а не предел архитектуры. Увеличение числа примеров даёт малый прирост: переход с 3000 на 10000 дал +1.6 пп token accuracy и +2.0 пп exact match.
• Узкий домен. Только синтетические склейки MNIST. Любая реальная фотография — вне обучающего распределения.
• Только 1–5 цифр. Другое количество цифр модель не обрабатывает.
• Нет устойчивости к деформациям. Ни поворотов, ни сдвигов, ни размытия при обучении не применялось (noise: 0.0).
• Валидация мала. 500 полосок дают стандартную ошибку ±1.6 пп при p=0.85, так что числа на 1–3 пп варьируются от запуска к запуску.
• Нет выделенного тестового сплита. Отчётные цифры — на валидации из MNIST test split.
• Только CPU-рецепт. Обучение и инференс на CPU; GPU-путь в конфиге есть, но не проверен и не использовался.

Model history

Текущий бандл — 10000 полосок/эпоху, 6 эпох, val_token_acc = 0.850.

Первая версия весов (3000 полосок) обучалась с утечкой метки и была непригодна. Исправлено три бага в model.py:

1. Утечка в forward. logits[:, Np:, :] — позиция Np содержит embedding истинного токена и предсказывала этот же токен, то есть модель училась копировать вход вместо чтения картинки. Правильный сдвиг: logits[:, Np - 1 : T_total - 1, :].
2. generate без причинной маски. Патчи видели друг друга двунаправленно, хотя при обучении маска была причинная.
3. generate прогонял выходы блоков через блоки повторно — torch.cat([cur, ...]), где cur уже результат всех 6 блоков, давал 12 слоёв вместо 6.

Симптом был обманчивым: train_acc показывал 0.98, а сгенерированные картинки читались как мусор. На честной метрике старые веса дают 0.037 token accuracy — уровень случайного угадывания.

Все бандлы после этих исправлений — только исправленный код. Валидация после фиксов:

Дополнительно было исправлено: потери в логе были завышены в grad_accum_steps раз; Ctrl+C теперь сохраняет бандл вместо потери прогресса.

Citation

Если пригодится в работе:

@misc{mnist_multidigit_gpt,
  title  = {mnist-multidigit-gpt: patch-based decoder-only digit sequence recognition},
  note   = {Custom PyTorch implementation, 4.8M params, trained on 10000 synthetic MNIST strips per epoch},
  year   = {2026},
  license = {GPL-3.0}
}

Обучающие данные:

@inproceedings{lecun1998mnist,
  title   = {Gradient-Based Learning Applied to Document Recognition},
  author  = {LeCun, Yann and Cortes, Corinna and Burges, Christopher J. C.},
  booktitle = {Proceedings of the IEEE},
  year    = {1998}
}

License

GPL-3.0 — на код и веса. Обучающие данные MNIST — CC BY-SA 3.0, атрибуция обязательна. Загружая эти веса, вы принимаете условия GPL-3.0 на производные работы.