Перенёс ByteTrack на GPU и ускорил мульти-камерный трекинг в 6 раз

Как я взял самый популярный трекинг-алгоритм, заставил его считать математику всех камер одним GPU-батчем — и по дороге трижды наступил на грабли, о которых стоит знать каждому, кто пишет real-time пайплайны на PyTorch.

Проблема: N камер = N трекеров = грустный CPU

Типичный продакшн-сценарий видеоаналитики: один сервер с GPU, на него приходит 8–32 RTSP-потока, на каждом кадре YOLO находит людей/машины, дальше нужен трекинг — стабильные ID объектов от кадра к кадру.

С детекцией всё хорошо: YOLO прекрасно батчится, 16 кадров прогоняются одним инференсом. А вот дальше начинается интересное. Все популярные Python-библиотеки трекинга — BoxMOT, Norfair, roboflow/trackers, встроенный трекер ultralytics — устроены одинаково: один экземпляр трекера на один поток видео, вся математика на CPU. Официальный рецепт для нескольких камер — «заверните каждую камеру в свой поток/процесс».

На 2–4 камерах это работает. На 16 — фильтр Кальмана, IoU-матрицы и управление треками съедают ядра быстрее, чем YOLO успевает детектировать: в наших замерах 16 потоков ByteTrack на CPU стоят 104 мс на кадр — против 6 мс на детекцию. Трекер, который «легче детектора», становится бутылочным горлышком пайплайна.

Единственное решение с настоящим кросс-камерным батчингом — NVIDIA DeepStream (nvtracker). Но это GStreamer-пайплайн, C+±плагины, YAML-конфиги и закрытое ядро трекера — совсем другая весовая категория, если весь ваш пайплайн на Python. (Подробное сравнение — в конце статьи.)

Захотелось третьего варианта: эффективность DeepStream, но pip install и три файла PyTorch. Так появился BatchGpuByteTrack.

Идея: трекер как батчевая операция

Ключевое наблюдение: математика ByteTrack по своей природе батчится не хуже нейросети.

• Kalman predict/update — умножения маленьких матриц (8×8). Тысяча треков — это тензор [1000, 8, 8], один вызов torch.einsum вместо тысячи вызовов numpy.

• IoU cost-матрицы — попарные пересечения боксов. Все потоки пакуются в один padded-тензор [S, M_max, N_max] — одно ядро.

• ReID-эмбеддинги — кропы всех детекций всех камер собираются в один батч для OSNet.

На CPU остаётся то, что действительно должно быть на CPU: венгерский алгоритм (lap.lapjv на матрицах «десятки×десятки» решается за микросекунды — GPU-реализации LAP окупаются на миллионах переменных, не на нашем размере) и управление жизненным циклом треков.

Получается конвейер: одна батчевая GPU-фаза (predict + IoU + OSNet) → одна выгрузка на CPU → параллельная ассоциация по потокам → одна батчевая GPU-фаза (update). Весь трафик CPU↔GPU — фиксированное число bulk-копий на кадр, не зависящее от числа камер.

При этом сам алгоритм — честный ByteTrack: три стадии ассоциации (включая фирменное «BYTE»-восстановление по низкоконфиденсным детекциям), стадия подтверждения новых треков, те же пороги. Паритет с референсной реализацией закреплён тестами.

Дальше начались грабли

Первая версия работала, но профайлер показывал странное: на 16 потоках GPU-версия обгоняла CPU всего в 1.2 раза. Разбор полётов вылился в три отдельные истории.

Грабли №1: «батчевый» код с циклом по потокам

Классика жанра. Kalman predict был уже батчевым — одна операция для всех треков всех камер. А Kalman update — нет:

# до: на каждый поток — свой upload, kernel и downloadfor stream_id in range(num_streams):    means, covs, meas = collect_matched(stream_id)          # CPU    means_t = torch.from_numpy(means).to(device)            # upload ×16    upd_means, upd_covs = kalman.batch_update(means_t, ...) # kernel ×16    results[stream_id] = upd_means.cpu().numpy()            # download ×16

16 потоков — 48 мелких операций с GPU на каждый кадр. Каждая — микросекунды полезной работы и десятки микросекунд накладных: запуск ядра, синхронизация, PCIe-транзакция.

Лечение тривиальное — собрать сматченные пары всех потоков в один тензор:

# после: один upload, один kernel, один download — для всех потоковentries, means, covs, meas = collect_matched_all_streams()upd_means, upd_covs = kalman.batch_update(stack(means), stack(covs), stack(meas))scatter_back(entries, upd_means.cpu().numpy())

Этап упал с 11.8 мс до 1.6 мс. Тот же приём для IoU-матриц (padded-батч [S, M_max, N_max] вместо цикла) снял ещё 5 мс:

Мораль: если в «GPU-ускоренном» коде есть for по потокам/объектам с .to(device) внутри — это не GPU-код, это генератор оверхеда.

Грабли №2: заливка кадров на 1.6 ГБ/с при шине в 25

Для ReID нужно доставить кадры на GPU. Микропрофилирование показало: заливка 16 кадров 720p (44 МБ) занимает 26.7 мс — это 1.6 ГБ/с при PCIe 4.0, умеющей ~25.

Причина — pageable-память: обычный torch.from_numpy(arr).to('cuda') из numpy-массива идёт через промежуточный staging-буфер драйвера. Лечение — свой персистентный pinned-буфер:

buf = torch.empty(F, H, W, 3, dtype=torch.uint8, pin_memory=True)  # один раз...for i, frame in enumerate(frames):                       # каждый кадр    buf[i].copy_(torch.from_numpy(frame))                # memcpy в pinnedimgs = buf.to(device, non_blocking=True)                 # DMA на полной скорости

Заодно кропы всех детекций стали вырезаться и ресайзиться одним вызовом torchvision.ops.roi_align по загруженным кадрам — вместо цикла «вырезали кроп → залили → отресайзили» на каждую детекцию (тот же антипаттерн, что в граблях №1, только с картинками).

Грабли №3: FP16 и 300 миллисекунд за каждый новый размер батча

Самое коварное. Включили FP16-инференс OSNet (эмбеддинги это переживают спокойно, скорость ×2) — и на тестовом прогоне появились необъяснимые фризы: медиана извлечения фич 14 мс, но регулярные кадры по 300 мс.

Оказалось: cuDNN подбирает ядра под конкретную форму входа, и под FP16 этот подбор стоит ~300 мс на каждый НОВЫЙ размер батча. А размер батча — это число детекций в кадре, оно меняется каждый кадр: 48 кропов, 52, 47, 61… Каждое новое число — новый фриз. На FP32 тот же прогрев стоит ~40 мс, поэтому там никто его не замечает.

Лечение — бакетирование: паддим батч инференса до кратного 32 (лишние строки отбрасываем после). Уникальных форм становится 3–4 на весь прогон, и каждая прогревается один раз — а можно прогреть все заранее при инициализации:

Если у вас в проде PyTorch-модель под autocast получает вход переменного размера — проверьте свои латентности на p99. Возможно, вы платите этот налог прямо сейчас.

Что получилось: цифры

Всё меряется двумя скриптами из репозитория, воспроизводимо на любом железе. Стенд: RTX 4090, PyTorch 2.13, 200 кадров, синтетическая нагрузка 30 детекций/поток/кадр (плотная сцена):

End-to-end с реальным YOLO на пешеходном видео (8 потоков): трекинг 5.1 мс против 31.8 мс на CPU (6.2x), суммарно с детекцией — 91 FPS на батч из 8 камер.

На 1–2 камерах CPU выигрывает: фиксированные накладные GPU там не окупаются.

ReID

Появление ReID меняет расклад: 70–80% кадра — это инференс OSNet, несжимаемый компьют. Батчевый roi_align + pinned-буферы + FP16 ускорили извлечение фич в ~2.3 раза:

Но чудес нет: с ReID общий выигрыш GPU над CPU — 1.8x на 8 потоках (закон Амдала: обе стороны платят за сеть сопоставимо, пока кропов мало). Дальше либо reid_interval (фичи раз в K кадров — треки живут на сглаженных эмбеддингах), либо TensorRT — он в роадмапе.

Что НЕ сработало (и почему это в репозитории)

• LAPJV-батчинг венгерки — параллелить решение крошечных матриц бессмысленно, оверхед тредов больше решения.

• torch.compile kernel fusion — динамические формы (число треков меняется каждый кадр) вызывают перекомпиляции, съедающие весь выигрыш. Вернёмся к этому после перевода состояния Кальмана в резидентные GPU-тензоры (тогда формы станут статическими).

• Венгерка на GPU — GPU-реализации LAP окупаются от миллионов переменных; наши матрицы — «десятки×десятки».

Все три тупика описаны в docs/optimization_guide.md.

Сравнение с NVIDIA DeepStream (nvtracker)

Единственный серьёзный аналог по идее кросс-камерного батчинга. Сравнение по ключевым параметрам:

Резюме: если вы уже живёте в экосистеме DeepStream или вам нужен NvDCF/Jetson-специфика — берите nvtracker, он взрослее. Если ваш пайплайн — Python/PyTorch, детектор свой, и вы хотите читать и менять код трекера — для этого и сделан BatchGpuByteTrack.

Как попробовать

git clone https://github.com/sweetlhare/BatchGpuByteTrack.gitcd BatchGpuByteTrack && pip install -r requirements.txtpython -m pytest tests/ -q          # 64 тестаpython tools/benchmark_comparison.py \    --video test_videos/6387-191695740_medium.mp4 --streams 1,4,8,16 --device cuda

from bytetrack import BatchGPUTracker, TrackerConfigtracker = BatchGPUTracker(num_streams=8, config=TrackerConfig(device='cuda'))tracks = tracker.update(batch_detections)   # список numpy [N,5] на каждый поток

Репозиторий: https://github.com/sweetlhare/BatchGpuByteTrack — тесты, бенчмарки, полная документация и роадмап оптимизаций внутри. Issues и PR приветствуются: особенно интересны замеры на других GPU и реальных мульти-камерных нагрузках.

---

Основано на ByteTrack (Zhang et al., ECCV 2022) и OSNet (Zhou et al., ICCV 2019), MIT-лицензия.