Обучение OCR-модели распознавания автомобильных номеров России на PyTorch

TL;DR

Задача — распознавание текста российских автомобильных номеров с изображения (кропа номерного знака).

Архитектура CRNN (CNN + BiLSTM + CTC) — классический и легковесный подход для OCR, ~9M параметров.

Датасет: 37 775 обучающих и 4 891 валидационное изображение. Алфавит — 22 символа (цифры + кириллица).

Mixed Precision (AMP) на NVIDIA RTX 5090 даёт ~2x ускорение обучения.

Итог: plate accuracy 95%+ после 50–100 эпох обучения с warmup и cosine annealing.

Постановка задачи

Распознавание текста российских автомобильных номеров с изображения — задача оптического распознавания символов (OCR), формулируемая как sequence-to-sequence прогноз: на входе — изображение номера фиксированного размера, на выходе — последовательность символов.

Где используется: в составе Video Analytics Platform для распознавания номеров в видеопотоке в реальном времени.

Архитектура модели

Изображение (1x32x256)
    ↓
[CNN Backbone] — Feature Extraction
    ↓
Tensor (B, 512, 1, 63) → (B, 63, 512)
    ↓
[BiLSTM x2] — Sequence Modeling
    ↓
[Linear 512→23] — Classification
    ↓
[CTC Decoding] — Greedy Decode
    ↓
Текст номера

CRNN (Shi et al., 2015) — классическая архитектура для OCR. Три блока: CNN извлекает признаки, BiLSTM моделирует последовательность, CTC обучает без посимвольной разметки.

CNN-бэкбон (Feature Extraction)

Свёрточная нейросеть извлекает из изображения последовательность признаков. Архитектура состоит из 5 блоков с BatchNorm после каждого свёрточного слоя.

Ключевые решения:

1) BatchNorm после каждого слоя — стабилизирует обучение и позволяет использовать более высокие learning rates.

2) MaxPool(2,1) в блоках 3 и 4 — уменьшает только высоту, сохраняя ширину (горизонтальное разрешение важно для последовательности символов).

3) Последний Conv с kernel=2 без padding — сводит высоту к 1, подготавливая данные для RNN.

Итоговый тензор (B, 512, 1, 63) переставляется в формат (B, 63, 512) — 63 временных шага, каждый с 512 признаками.

BiLSTM (Sequence Modeling)

Двунаправленная LSTM обрабатывает последовательность признаков, захватывая контекст в обоих направлениях. Это критично для понимания границ символов.

nn.LSTM(
    input_size=512,
    hidden_size=256,
    num_layers=2,
    bidirectional=True,
    dropout=0.2,
)

Конфигурация: input_size=512, hidden_size=256, 2 слоя, bidirectional=True, dropout=0.2.

Bidirectional означает, что каждый временной шаг получает информацию как слева, так и справа. 2 слоя дают достаточную глубину, а dropout 0.2 — регуляризацию.

CTC Loss и декодирование

CTC (Connectionist Temporal Classification) — ключевой компонент для обучения без посимвольной разметки. Мы знаем только итоговый текст номера, но не знаем, какой пиксель соответствует какому символу.

Вводится blank-токен (индекс 0) — «нет символа». Модель предсказывает распределение по всем 23 классам (22 символа + blank) для каждого из 63 временных шагов.

CTC loss суммирует вероятности всех путей выравнивания, которые приводят к целевому тексту. zero_infinity=True — защита от NaN на ранних этапах.

Декодирование (greedy): для каждого временного шага берётся символ с максимальной вероятностью, затем удаляются дубликаты и blank-токены.

Подготовка данных

Датасет: 37 775 обучающих и 4 891 валидационное изображение (~88.5% / 11.5%). Лейблы извлекаются из имён файлов (например, A001AA50.png → A001AA50).

Предобработка: OpenCV (BGR→RGB), grayscale (1 канал), ресайз до 32×256, нормализация к диапазону [-1, 1].

Аугментации (только для train)

Для повышения устойчивости к реальным условиям съёмки:

Процесс обучения

Оптимизатор Adam с learning rate 1e-4. Warmup + Cosine Annealing: первые 5 эпох LR линейно возрастает от 0 до 1e-4 (плавный старт), затем косинусное затухание до 1e-12.

Mixed Precision (AMP): forward pass в float16, GradScaler предотвращает underflow градиентов. На RTX 5090 — ~2x ускорение и увеличение batch size.

Gradient Clipping по норме max_norm=5.0 — предотвращает «взрыв градиентов» в LSTM, особенно на ранних этапах обучения.

Модель сохраняется при улучшении validation loss (не train) — предотвращает переобучение.

Гиперпараметры

Метрики и результаты

Сравнение с альтернативами

В ходе работы тестировались два подхода:

Почему PyTorch: TensorFlow с версии 2.15 прекратил поддержку GPU на Windows. PyTorch с CUDA 12.8 полностью использует RTX 5090 — ускорение обучения в десятки раз по сравнению с CPU.

Технологический стек

Типичные ошибки при обучении OCR

1) Обучать без аугментаций — модель не переносится на реальные условия (шум, освещение, blur).

2) Использовать RGB вместо grayscale — номера не требуют цветовой информации, а 3 канала увеличивают модель.

3) Не использовать gradient clipping — LSTM склонна к «взрыву градиентов» на ранних этапах.

4) Сохранять модель по train loss, а не по val loss — приводит к переобучению.

5) Начинать с полного LR без warmup — деструктивные обновления весов на раннем этапе.

6) Игнорировать формат номеров (разная длина) — без правильной collate-функции CTC не обучится.

Мониторинг обучения

Интеграция с TensorBoard для визуализации: графики train/val loss по эпохам, plate accuracy и character accuracy, изменение learning rate.

Запуск: tensorboard --logdir runs/v2_ocr

FAQ

Key Takeaways

1) CRNN + CTC — эффективная и легковесная архитектура для OCR номеров (~9M параметров).

2) Mixed Precision + RTX 5090 дают значительное ускорение обучения.

3) Warmup + Cosine Annealing обеспечивают стабильную сходимость без осцилляций.

4) Аугментации критичны для устойчивости к шуму и освещению в реальных условиях.

5) PyTorch — единственный вариант для GPU-обучения на Windows на текущий момент.