From 89ca5b41633d81653f31d4a529ab61423cfbeae2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=BB=20=D0=9E=D0=BC=D0=B5=D0=BB?= =?UTF-8?q?=D0=B5=D1=87=D0=BA=D0=BE?= Date: Sat, 26 Sep 2026 04:10:04 +0300 Subject: [PATCH] =?UTF-8?q?=D0=BF=D0=BB=D0=BE=D1=82=D0=BD=D1=8B=D0=B5=20?= =?UTF-8?q?=D1=81=D1=82=D0=B0=D0=B4=D0=B8=D0=B8=20=D0=BD=D0=B0=20=D0=B2?= =?UTF-8?q?=D0=B8=D0=B4=D0=B5=D0=BE=D0=BA=D0=B0=D1=80=D1=82=D0=B5,=20?= =?UTF-8?q?=D0=BE=D1=82=D0=BA=D0=B0=D1=82=20=D0=BD=D0=B0=20=D0=BF=D1=80?= =?UTF-8?q?=D0=BE=D1=86=D0=B5=D1=81=D1=81=D0=BE=D1=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 +- docs/EXPERIMENTS.md | 95 ++++++++ flyguard/gpu.py | 500 +++++++++++++++++++++++++++++++++++++++++ flyguard/lobula.py | 6 +- flyguard/pipeline.py | 67 ++++-- requirements.txt | 4 +- tests/test_pipeline.py | 107 +++++++++ 7 files changed, 763 insertions(+), 18 deletions(-) create mode 100644 flyguard/gpu.py diff --git a/README.md b/README.md index 526af91..6afe43f 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ flyguard/ ядро: стадии обработки, память, сч export.py 3D-рамки, время до столкновения, маркеры RViz synth.py вставка предметов трассировкой лучей tools/ обучение, оценка, разбор, полигон с аугментациями -tests/ 47 тестов, запускаются без данных и без ROS +tests/ 50 тестов, запускаются без данных и без ROS docs/ методика и результаты artifacts/ обученные модели ``` diff --git a/docs/EXPERIMENTS.md b/docs/EXPERIMENTS.md index 56abcf5..7dcf7ee 100644 --- a/docs/EXPERIMENTS.md +++ b/docs/EXPERIMENTS.md @@ -3223,3 +3223,98 @@ MBON, новизна; у каждого трека ход улики) на си | контур габарита в кривой | прямой короб, изогнутой части не видно | оба контура | | решения узла | — | без изменений (синтетика: те же находки, фантомы, дальности и число кадров) | | тесты | 48 | 50 | + +## 21. Плотные стадии на видеокарте с откатом на процессор (26.09) + +Решение пользователя: считать на видеокарте, если она есть, и только при её +отсутствии или сбое — на процессоре. Раньше (п. 7.4) вывод держался на +процессоре: на видеокарту была перенесена одна ламина, выигрыш 3 мс из 43. +Перенос одной стадии и правда ничего не даёт — выигрыш даёт перенос всех +плотных стадий разом, с данными, которые между стадиями не ездят. + +### 21.1. Что переносить + +Профиль кадра (cProfile, синтетика, 150 кадров; под профилировщиком всё +медленнее, важны доли): кластеризация с учётом глубины 20 мс, раскладка точек +по углам 14, ламина 8, оценка движения 8.5, ось пути 6.5, плоскость 3.9 из 74. +Без профилировщика кадр — 49 мс на синтетике и 36 мс на записи с поезда, из них +на раскладку, ламину и кластеризацию приходится 22–34 мс. Они работают с целым +образом 128 × 600 лучей — это и перенесено (`flyguard/gpu.py`). Плоскость, ось, +движение, треки и решение остались на процессоре: там мелкие массивы и +ветвистая логика, копирование на видеокарту дороже расчёта. + +### 21.2. Те же операции в тех же типах + +Цель — не «похожий», а тот же результат, иначе откат на процессор менял бы +решения посреди поездки. + +* **Сетчатка** — выборки, умножения и корень в float32, отдельными ядрами + (без слияния умножения со сложением). На видеокарту уходит только окно сырых + столбцов, нужное сектору (5 МБ вместо 24). +* **Ламина.** `scipy.ndimage.uniform_filter` считает окно скользящим средним в + double. Суммы чисел float32 такого диапазона (диспаритет 0.003…20, окно до + 121) в double точны при любом порядке сложения, поэтому окно берётся из + накопленной суммы, делится на размер и округляется в float32 после каждой + оси — как в scipy. Сверка окон: 0 расхождений на 24.5 млн значений (запись) и + 20.6 млн (синтетика). Вся ламина записана в граф CUDA: видеокарта считает её + за 0.4 мс, а запуск сотни мелких операций из Python стоил 3 мс; с графом — + 0.6 мс вместе с выгрузкой. +* **Кластеризация** — рёбра сразу для всех 17 соседей, компоненты + подвешиванием к меньшему номеру со сжатием путей: корень компоненты — её + наименьший луч, и нумерация совпадает со `scipy.sparse.csgraph`. На 42 тыс. + лучей и 581 тыс. рёбер — 8 раундов (медиана). + +Исключение — раскладка по углам точек (кадры со сбитым порядком, как в +синтетике организаторов): арктангенс видеокарты и процессора расходится в +последнем знаке, а numpy сортирует ячейки неустойчиво. + +### 21.3. Сверка + +По стадиям, 80 кадров: на `roundT_doubleT` все три — бит в бит; на синтетике +ламина и кластеризация — бит в бит, раскладка по углам разошлась в 13 кадрах из +80, 326 значений из 300 тыс. + +Конвейер целиком, синтетика и шесть записей, покадрово против процессорного +прогона той же конфигурации (`v_t33`): на всех шести записях совпали решения, +объекты, кандидаты и треки в каждом кадре; на синтетике решения, объекты и +треки — во всех 1499 кадрах, кандидаты разошлись в трёх (раскладка по углам), +на треки это не повлияло. Синтетика 9 из 10, те же дальности и число кадров. + +### 21.4. Откат + +Любое исключение видеокарты — нет драйвера, не хватило памяти, ошибка ядра — +ловит конвейер (`FlyGuard._gpu_off`): до конца работы он считает на +процессоре, а текущий кадр досчитывается там же. Тест +`test_pipeline_survives_gpu_failure_mid_frame` роняет видеокарту на 20-м кадре +и сверяет решения с чисто процессорным прогоном. Узел пишет в журнал, где идёт +счёт, и предупреждает, если видеокарта отказала на ходу; в диагностике — поле +`device`. + +### 21.5. В контейнере + +Образ с PyTorch 2.9.1 + CUDA 12.8, Docker в WSL, видеокарта через `/dev/dxg` +(на стенде — `--gpus all`). Узел, `ros2 bag play` в том же контейнере: + +| запись | счёт | кадр, медиана / p95 | обнаружение | +|---|---|---|---| +| `doubleT_obstacle` | видеокарта | 23.5 / 31.3 мс | 189 из 190, 55.8 м | +| `doubleT_obstacle` | процессор | 32.3 / 36.2 мс | 189 из 190, 55.8 м | +| синтетика | видеокарта | 26.8 / 43.1 мс | 390 кадров с тревогой | +| синтетика | процессор | 42.3 / 47.0 мс | 398 кадров с тревогой | + +Строка `doubleT_obstacle` с видеокартой — после правки прогрева. До неё кадров +с видеокартой приходило меньше (195 и 1502 против 201 и 1510, отсюда и 390 +кадров тревоги на синтетике против 398): первый кадр шёл 1.0–1.7 с — при первом +вызове подгружаются ядра CUDA раскладки, и очередь подписки переполнялась. +Прогрев (`GpuStages.warmup`) теперь гоняет и раскладку, на маленькой выдуманной +решётке обоими путями: принято 201 из 201, худший кадр 38 мс. + +### 21.6. Образ + +PyTorch с библиотеками CUDA — 6.6 ГБ (сами библиотеки NVIDIA 4.3 ГБ, cuDNN и +cuBLAS из них почти половина); образ — 9.3 ГБ вместо 2.7. Нашим стадиям cuDNN, +cuBLAS и прочие не нужны, но PyTorch без них не загружается. Ставится одним +слоем без кэша pip. Колёса — из индекса PyTorch или, где сети нет, с HTTP-адреса +`TORCH_WHEELS` (`docker/fetch_wheels.py`, `docker/wheels/README.md`). Через +контекст сборки колёса не передаются: первая попытка так делала, и контекст в +4.1 ГБ плюс слой с колёсами плюс установка чуть не заняли весь диск C:. diff --git a/flyguard/gpu.py b/flyguard/gpu.py new file mode 100644 index 0000000..70fdf38 --- /dev/null +++ b/flyguard/gpu.py @@ -0,0 +1,500 @@ +"""Видеокарта для плотных стадий конвейера: сетчатка, ламина, кластеризация. + +Три стадии работают с целым образом 128 × 600 лучей и занимают больше +половины кадра: раскладка точек по решётке, ламина (центр минус окружение на +трёх масштабах) и кластеризация лучей с учётом глубины. Они и перенесены. +Остальное — плоскость пути, ось, оценка движения, треки, решение — остаётся +на процессоре: там мелкие массивы и ветвистая логика, и копирование на +видеокарту стоило бы дороже самого расчёта. + +Каждая функция повторяет процессорную версию операция в операцию, в тех же +типах. Раскладка по порядку точек, ламина и кластеризация дают результат, +совпадающий с процессорным бит в бит (тесты и сверка на записях, +EXPERIMENTS п. 21): + +* сетчатка — только выборки, умножения и корень, всё в float32 без слияния + операций; +* ламина — равномерные окна считаются через накопленные суммы в float64 и + округляются в float32 после каждой оси, ровно как `scipy.ndimage. + uniform_filter`: сумма чисел float32 такого диапазона в float64 точна, и + результат не зависит от порядка сложения; +* кластеризация — целочисленный граф; компоненты нумеруются по наименьшему + лучу, как у `scipy.sparse.csgraph.connected_components`. + +Исключение — раскладка по углам точек (кадры с нарушенным порядком, как в +синтетике организаторов): арктангенс видеокарты и процессора может разойтись +в последнем знаке, а numpy сортирует ячейки неустойчиво. Расходятся считанные +лучи на кадр. + +Любой сбой видеокарты — нет драйвера, не хватило памяти, ошибка ядра — ловит +конвейер: до конца работы он считает на процессоре (`device. +notify_cuda_error`), а кадр досчитывается там же, решение не теряется. +""" +from __future__ import annotations + +import warnings + +import numpy as np + +from .lamina import SCALES, LaminaOutput +from .retina import DEG, ORDER_TOL_DEG, RangeImage, ScanLayout + + +class GpuStages: + """Плотные стадии на видеокарте; держит таблицы решётки и образ кадра.""" + + def __init__(self, device: str = "cuda"): + import torch + self.torch = torch + self.dev = torch.device(device) + torch.zeros(1, device=self.dev) # поднять контекст сразу, а не на первом кадре + self._tab_key = None + self._tab: dict = {} + self._win: dict = {} + self._off: dict = {} + self._box: dict = {} + self._graphs = True # графы CUDA для ламины; при отказе — по одной операции + self._lam_key = None + # образ текущего кадра: ламина и кластеризация берут его отсюда, не копируя заново + self.r = None + self.valid = None + self._r_host = None + + def warmup(self, shape: tuple[int, int] = (128, 600)) -> None: + """Прогнать все три стадии на выдуманном кадре. + + Первый запуск каждого ядра CUDA подгружает его код: без прогрева первый + кадр записи шёл 1–1.7 с, очередь подписки переполнялась, и терялись + первые 6–8 кадров (замерено в контейнере). Раскладка прогревается на + маленькой решётке обоими путями — по порядку точек и по углам. + """ + from .cdr import PointCloud2 + t = self.torch + rng = np.random.default_rng(0) + n, w, e = 8, 64, 2 + lay = ScanLayout(np.linspace(10.0, -20.0, n), -0.1, 3.2, np.zeros(n, np.int64), + np.zeros(n), w, e) + xyz = lay.dirs.transpose(1, 0, 2)[:, None, :, :] * np.float32(20.0) # столбец·эхо·кольцо + xyz = np.broadcast_to(xyz, (w, e, n, 3)).reshape(-1, 3) + pts = np.zeros(xyz.shape[0], dtype=[("x", " dict: + key = (id(L), L.n_points, L.n_az) + if key != self._tab_key: + t, d = self.torch, self.dev + + def up(a, dtype): + return t.as_tensor(np.ascontiguousarray(a), dtype=dtype, device=d) + + self._tab = { + "sin_el": up(L._sin_el, t.float32), + "el_asc": up(L._el_asc, t.float64), + "el_order": up(L._el_order, t.int64), + "el_deg": up(L.el_deg, t.float64), + "el_step": up(L.el_step_deg, t.float64), + "resid": up(L.az_resid_deg, t.float64), + } + self._win = {} + self._tab_key = key + return self._tab + + def _window(self, L: ScanLayout, start: int, stop: int) -> tuple: + """Сырые столбцы, нужные сектору, и карта выпрямления внутри них.""" + key = (start, stop) + if key not in self._win: + w = L.n_az + lo = start + int(L.col_shift.min()) + hi = stop + int(L.col_shift.max()) + g = L.gather[:, start:stop] + ok = L.gather_ok[:, start:stop] + if L.wrap: + raw = np.arange(lo, hi) % w + g = (g - lo) % w + else: + lo, hi = max(lo, 0), min(hi, w) + raw = np.arange(lo, hi) + g = g - lo + ok = ok & (g >= 0) & (g < (hi - lo)) + g = np.clip(g, 0, hi - lo - 1) + # столбцы окна идут подряд, кроме стыка кругового скана + cut = np.flatnonzero(np.diff(raw) != 1) + 1 + runs = [(int(a[0]), int(a[-1]) + 1) for a in np.split(raw, cut) if a.size] + t = self.torch + self._win[key] = (runs, t.as_tensor(g, dtype=t.int64, device=self.dev), + t.as_tensor(ok, device=self.dev)) + return self._win[key] + + # ------------------------------------------------------------------ загрузка точек + + def _upload(self, pts: np.ndarray, rows: list[tuple[int, int]]): + """Строки структурированного массива → байты на видеокарте, (N, шаг точки).""" + t = self.torch + step = pts.dtype.itemsize + parts = [] + for a, b in rows: + raw = np.ascontiguousarray(pts[a:b]).view(np.uint8) + with warnings.catch_warnings(): + # буфер сообщения только для чтения, а тензор мы не пишем + warnings.simplefilter("ignore", UserWarning) + parts.append(t.from_numpy(raw).to(self.dev, non_blocking=False)) + buf = parts[0] if len(parts) == 1 else t.cat(parts) + return buf.view(-1, step) + + def _field(self, buf, pts: np.ndarray, name: str): + dt, off = pts.dtype.fields[name][:2] + if dt.kind != "f" or dt.itemsize != 4 or dt.byteorder == ">": + raise TypeError(f"поле {name}: ожидался float32, а не {dt}") + return buf[:, off:off + 4].contiguous().view(self.torch.float32).reshape(-1) + + # ------------------------------------------------------------------ сетчатка + + def project(self, L: ScanLayout, pc, cols: slice | None) -> RangeImage: + """То же, что `ScanLayout.project`, на видеокарте.""" + w = L.n_az + start = 0 if cols is None else (cols.start or 0) + stop = w if cols is None else (cols.stop if cols.stop is not None else w) + self._tables(L) + img = None + if L.indexed and pc.n_points == L.n_points: + img = self._project_indexed(L, pc, start, stop) + if img is None: + L.n_geometric += 1 + img = self._project_geometric(L, pc, start, stop) + self._r_host = img.r_near + return img + + def _project_indexed(self, L: ScanLayout, pc, start: int, stop: int): + t = self.torch + tab = self._tab + n, e = L.n_rings, L.n_echo + runs, g, ok = self._window(L, start, stop) + pts = pc.points + per_col = e * n + buf = self._upload(pts, [(a * per_col, b * per_col) for a, b in runs]) + ncol = sum(b - a for a, b in runs) + + def cube(name: str): + return self._field(buf, pts, name).reshape(ncol, e, n).permute(2, 0, 1) + + x, y, z = cube("x"), cube("y"), cube("z") + good = (x != 0) | (y != 0) | (z != 0) + r = t.sqrt(x * x + y * y + z * z) + good &= t.isfinite(r) + r = t.where(good, r, t.zeros((), dtype=r.dtype, device=self.dev)) + + dev = t.abs(z - r * tab["sin_el"][:, None, None]) + tol = r * np.float32(ORDER_TOL_DEG * DEG) + np.float32(1e-3) + if bool(t.any(dev > tol)): + return None + + if e == 1: + r_near = r[..., 0] + r_far = r[..., 0] + it = cube("intensity")[..., 0] + valid = good[..., 0] + else: + inten = cube("intensity") + inf = t.full((), float("inf"), dtype=r.dtype, device=self.dev) + near_i = t.argmin(t.where(good, r, inf), dim=-1, keepdim=True) + far_i = t.argmax(r, dim=-1, keepdim=True) + r_near = t.gather(r, -1, near_i)[..., 0] + r_far = t.gather(r, -1, far_i)[..., 0] + it = t.gather(inten, -1, near_i)[..., 0] + valid = good.any(dim=-1) + + r_near = t.gather(r_near, 1, g) + r_far = t.gather(r_far, 1, g) + it = t.gather(it, 1, g) + valid = t.gather(valid, 1, g) & ok + zero = t.zeros((), dtype=r_near.dtype, device=self.dev) + r_near = t.where(valid, r_near, zero) + r_far = t.where(valid, r_far, zero) + return self._finish(pc.stamp, r_near, r_far, it, valid) + + def _project_geometric(self, L: ScanLayout, pc, start: int, stop: int): + t = self.torch + tab = self._tab + n, wid = L.n_rings, stop - start + pts = pc.points + buf = self._upload(pts, [(0, pts.shape[0])]) + x, y, z = (self._field(buf, pts, k) for k in ("x", "y", "z")) + good = (((x != 0) | (y != 0) | (z != 0)) & t.isfinite(x) & t.isfinite(y) + & t.isfinite(z)) + idx = t.nonzero(good, as_tuple=True)[0] + x, y, z = x[idx], y[idx], z[idx] + + step = L.az_step_deg + az = t.rad2deg(t.atan2(x, -y)) + jf = (az - np.float32(L.az0_deg)) / np.float32(step) + if L.wrap: + jf = t.remainder(jf, L.n_az) + margin = float(np.abs(L.az_resid_deg).max()) / abs(step) + 1.0 + sel = t.nonzero((jf > start - margin) & (jf < stop - 1 + margin), as_tuple=True)[0] + idx, x, y, z, az = idx[sel], x[sel], y[sel], z[sel], az[sel] + + r = t.sqrt(x * x + y * y + z * z) + el = t.rad2deg(t.asin(t.clamp(z / t.clamp_min(r, np.float32(1e-6)), -1.0, 1.0))) + el64 = el.double() + asc = tab["el_asc"] + k = t.clamp(t.searchsorted(asc, el64), 1, n - 1) + k = k - ((el64 - asc[k - 1]) < (asc[k] - el64)).long() + h = tab["el_order"][k] + ok = t.abs(el64 - tab["el_deg"][h]) <= t.clamp_min(tab["el_step"][h], 0.2) + j = t.round(((az.double() - tab["resid"][h]) - L.az0_deg) / step).long() + if L.wrap: + j = t.remainder(j, L.n_az) + ok &= (j >= start) & (j < stop) + sel = t.nonzero(ok, as_tuple=True)[0] + + r_near = t.zeros(n * wid, dtype=t.float32, device=self.dev) + r_far = t.zeros_like(r_near) + it = t.zeros_like(r_near) + valid = t.zeros(n * wid, dtype=t.bool, device=self.dev) + if sel.numel(): + cell = h[sel] * wid + (j[sel] - start) + rr = r[sel] + mm = t.clamp_max(rr * 1000.0, float((1 << 20) - 1)).long() + order = t.argsort(cell * (1 << 20) + mm, stable=True) + cs, rs = cell[order], rr[order] + new = t.ones(cs.numel(), dtype=t.bool, device=self.dev) + new[1:] = cs[1:] != cs[:-1] + first = t.nonzero(new, as_tuple=True)[0] + last = t.cat([first[1:] - 1, first.new_tensor([cs.numel() - 1])]) + r_near[cs[first]] = rs[first] + r_far[cs[last]] = rs[last] + valid[cs[first]] = True + if "intensity" in pts.dtype.names: + src = self._field(buf, pts, "intensity")[idx[sel]][order] + it[cs[first]] = src[first] + return self._finish(pc.stamp, r_near.reshape(n, wid), r_far.reshape(n, wid), + it.reshape(n, wid), valid.reshape(n, wid)) + + def _finish(self, stamp: float, r_near, r_far, it, valid) -> RangeImage: + self.r, self.valid = r_near, valid + host = [a.cpu().numpy() for a in (r_near, r_far, it, valid)] + return RangeImage(stamp, *host) + + # ------------------------------------------------------------------ ламина + + def _box_axis(self, X, sizes: list[int], dim: int): + """Равномерное окно по одной оси для пачки размеров сразу. + + `X` — (K, B, H, W) float32, k-й срез фильтруется окном `sizes[k]` + вдоль `dim` (2 или 3) с продолжением краевым значением. Суммы окон + берутся из накопленной суммы в float64: все частичные суммы чисел + float32 такого диапазона в float64 точны, поэтому сумма окна точна при + любом порядке сложения. Деление на размер и округление в float32 — + как в `scipy.ndimage.uniform_filter1d`. + """ + t = self.torch + Xd = X.double() + n = Xd.shape[dim] + lo, hi, n_lo, n_hi, size = self._box_tables(tuple(sizes), n, dim) + cs = t.cumsum(Xd, dim=dim) + cs = t.cat([t.zeros_like(cs.narrow(dim, 0, 1)), cs], dim=dim) + + def along(a): + return a.expand(*Xd.shape[:dim], n, *Xd.shape[dim + 1:]) + + s = (t.gather(cs, dim, along(hi)) - t.gather(cs, dim, along(lo)) + + along(n_lo) * Xd.narrow(dim, 0, 1) + along(n_hi) * Xd.narrow(dim, n - 1, 1)) + return (s / size).float() + + def _box_tables(self, sizes: tuple, n: int, dim: int): + """Границы окон по оси: считаются один раз (и до записи графа CUDA).""" + key = (sizes, n, dim) + if key not in self._box: + t = self.torch + p = t.as_tensor([s // 2 for s in sizes], dtype=t.int64, device=self.dev) + i = t.arange(n, device=self.dev) + shape = [len(sizes), 1, 1, 1] + shape[dim] = n + lo = (i[None, :] - p[:, None]).clamp(min=0).reshape(shape) + hi = ((i[None, :] + p[:, None]).clamp(max=n - 1) + 1).reshape(shape) + n_lo = (p[:, None] - i[None, :]).clamp(min=0).double().reshape(shape) + n_hi = (i[None, :] + p[:, None] - (n - 1)).clamp(min=0).double().reshape(shape) + size = t.as_tensor(sizes, dtype=t.float64, device=self.dev).reshape(-1, 1, 1, 1) + self._box[key] = (lo, hi, n_lo, n_hi, size) + return self._box[key] + + def _boxes(self, X, sizes: list[tuple[int, int]]): + """`scipy.ndimage.uniform_filter(x, size, mode="nearest")` для пачки. + + `X` — (B, H, W) float32; выход (len(sizes), B, H, W): по осям по + очереди, с округлением в float32 после каждой, как в scipy. + """ + K = len(sizes) + Y = self._box_axis(X[None].expand(K, *X.shape), [s[0] for s in sizes], 2) + return self._box_axis(Y, [s[1] for s in sizes], 3) + + def lamina(self, r_max: float = 300.0) -> LaminaOutput: + """То же, что `lamina._process_cpu`, по образу текущего кадра. + + Сама видеокарта считает ламину за доли миллисекунды, а запуск сотни + мелких операций из Python стоит в несколько раз дороже. Поэтому ламина + записывается в граф CUDA один раз на размер образа и дальше + запускается одной командой. Не поддерживает граф драйвер — считаем + теми же операциями по одной. + """ + t = self.torch + key = (tuple(self.r.shape), float(r_max)) + out = None + if self._graphs: + try: + if self._lam_key != key: + self._capture_lamina(key, r_max) + self._lam_in[0].copy_(self.r) + self._lam_in[1].copy_(self.valid) + self._lam_graph.replay() + out = self._lam_out + except Exception: + self._graphs = False + self._lam_key = None + if out is None: + out = self._lamina_core(self.r, self.valid, r_max) + f = out[0].cpu().numpy() + return LaminaOutput(disp=f[0], on=f[1], off=f[2], on_scale=out[1].cpu().numpy(), + surround=f[3], hole=f[4]) + + def _capture_lamina(self, key: tuple, r_max: float) -> None: + t = self.torch + self._lam_in = (t.empty_like(self.r), t.empty_like(self.valid)) + self._lam_in[0].copy_(self.r) + self._lam_in[1].copy_(self.valid) + side = t.cuda.Stream(self.dev) + side.wait_stream(t.cuda.current_stream(self.dev)) + with t.cuda.stream(side): + for _ in range(2): # прогрев: таблицы окон и кэш памяти + self._lamina_core(*self._lam_in, r_max) + t.cuda.current_stream(self.dev).wait_stream(side) + g = t.cuda.CUDAGraph() + with t.cuda.graph(g): + self._lam_out = self._lamina_core(*self._lam_in, r_max) + self._lam_graph, self._lam_key = g, key + + def _lamina_core(self, r, valid, r_max: float): + """Ламина целиком на видеокарте: (disp, on, off, surround, hole) и on_scale.""" + t = self.torch + v = valid.float() + zero = t.zeros((), dtype=t.float32, device=self.dev) + disp = t.where(valid & (r > 0.05), 1.0 / r, zero) * v + + # все окна кадра одной пачкой: на каждом масштабе центр и окружение, + # для диспаритета и для маски эха, плюс окно «дыры» + sizes = [] + for r_in, r_out in SCALES: + sizes += [(2 * r_in + 1, 2 * r_in + 1), (2 * r_out + 1, 4 * r_out + 1)] + sizes.append((5, 15)) + box = self._boxes(t.stack([disp, v]), sizes) + + on = t.zeros_like(disp) + off = t.zeros_like(disp) + on_scale = t.zeros(disp.shape, dtype=t.int8, device=self.dev) + surround_mid = None + for k, (r_in, r_out) in enumerate(SCALES): + s_in, s_out = sizes[2 * k], sizes[2 * k + 1] + n_in = s_in[0] * s_in[1] + n_out = s_out[0] * s_out[1] + num = box[2 * k + 1, 0] * n_out - box[2 * k, 0] * n_in + den = box[2 * k + 1, 1] * n_out - box[2 * k, 1] * n_in + sur = t.where(den > 0.5, num / den, zero) + enough = den > 8.0 + c = t.where(enough, disp - sur, zero) + pos = t.clamp_min(c, 0.0) * v + neg = t.clamp_min(-(disp - sur), 0.0) * enough + better = pos > on + on = t.where(better, pos, on) + on_scale = t.where(better, t.full((), k, dtype=t.int8, device=self.dev), on_scale) + off = t.maximum(off, neg) + if k == 1: + surround_mid = sur + hole = 1.0 - box[-1, 1] + on = t.clamp(on, 0.0, 1.0 / max(r_max, 1.0) * 1e4) + return t.stack([disp, on, off, surround_mid, hole]), on_scale + + # ------------------------------------------------------------------ кластеризация + + def _offsets(self, col_reach: int, row_reach: int): + key = (col_reach, row_reach) + if key not in self._off: + pairs = [(dr, dc) for dr in range(0, row_reach + 1) + for dc in range(-col_reach, col_reach + 1) if not (dr == 0 and dc <= 0)] + p = self.torch.tensor(pairs, dtype=self.torch.int64, device=self.dev) + self._off[key] = (p[:, 0:1], p[:, 1:2]) + return self._off[key] + + def cluster_by_depth(self, mask: np.ndarray, r: np.ndarray, *, + rel_tol: float = 0.06, abs_tol: float = 0.35, + col_reach: int = 3, row_reach: int = 2): + """То же, что `lobula.cluster_by_depth`, с теми же номерами компонент. + + Рёбра строятся сразу для всех соседей, компоненты — подвешиванием к + меньшему номеру со сжатием путей: у каждой компоненты корнем остаётся + её наименьший луч, и нумерация по корням совпадает со scipy. + """ + t = self.torch + h, w = mask.shape + if r is self._r_host and self.r is not None: + rt = self.r.reshape(-1) + else: + rt = t.as_tensor(np.ascontiguousarray(r), dtype=t.float32, + device=self.dev).reshape(-1) + m = t.as_tensor(np.ascontiguousarray(mask), device=self.dev).reshape(-1) + idx = t.nonzero(m, as_tuple=True)[0] + n = idx.numel() + if n == 0: + return np.zeros(mask.shape, np.int32), 0 + lut = t.full((h * w,), -1, dtype=t.int64, device=self.dev) + lut[idx] = t.arange(n, device=self.dev) + rows = t.div(idx, w, rounding_mode="floor") + cols = idx - rows * w + ri = rt[idx] + dr, dc = self._offsets(col_reach, row_reach) + nr = rows[None, :] + dr + nc = cols[None, :] + dc + ok = (nr < h) & (nc >= 0) & (nc < w) + nb = lut[t.where(ok, nr * w + nc, t.zeros_like(nr))] + ok &= nb >= 0 + rb = ri[nb.clamp(min=0)] + ra = ri[None, :].expand_as(rb) + ok &= t.abs(ra - rb) <= (abs_tol + rel_tol * t.minimum(ra, rb)) + s = t.arange(n, device=self.dev)[None, :].expand_as(nb)[ok] + d = nb[ok] + + parent = t.arange(n, device=self.dev) + while True: + for _ in range(4): + ps, pd = parent[s], parent[d] + parent.scatter_reduce_(0, t.maximum(ps, pd), t.minimum(ps, pd), reduce="amin") + parent = parent[parent] + parent = parent[parent] + if not bool((parent[s] != parent[d]).any()): + break + while True: + nxt = parent[parent] + if bool(t.equal(nxt, parent)): + break + parent = nxt + roots, comp = t.unique(parent, return_inverse=True) + labels = t.zeros(h * w, dtype=t.int32, device=self.dev) + labels[idx] = (comp + 1).to(t.int32) + return labels.reshape(h, w).cpu().numpy(), int(roots.numel()) diff --git a/flyguard/lobula.py b/flyguard/lobula.py index d6f7198..39d2522 100644 --- a/flyguard/lobula.py +++ b/flyguard/lobula.py @@ -292,7 +292,7 @@ def find_candidates(tf: TrackFrame, lam: LaminaOutput, corridor: Corridor, *, platform_h_hi: float = 1.25, h_top: float = 0.0, half_width_top: float = 0.0, top_d_max: float = 90.0, top_detect: bool = False, - top_min_rays: int = 3) -> list[Candidate]: + top_min_rays: int = 3, cluster=None) -> list[Candidate]: """Выделить кандидатов: связные пятна лучей, попавших в габарит. Кластеризация ведётся по **расширенной** области, а членство в габарите @@ -337,7 +337,9 @@ def find_candidates(tf: TrackFrame, lam: LaminaOutput, corridor: Corridor, *, gap_img = contrast_to_depth_gap(lam.on, tf.r) - labels, n = cluster_by_depth(context, tf.r, col_reach=dilate_cols + 1) + # `cluster` — та же кластеризация на видеокарте (FlyGuard._cluster), с тем + # же результатом; здесь самая большая маска кадра, и выигрыш в ней + labels, n = (cluster or cluster_by_depth)(context, tf.r, col_reach=dilate_cols + 1) if n == 0: return [] if split_gap > 0: diff --git a/flyguard/pipeline.py b/flyguard/pipeline.py index 0589611..d52c0d7 100644 --- a/flyguard/pipeline.py +++ b/flyguard/pipeline.py @@ -22,7 +22,7 @@ from .descending import Decision, DescendingNeurons from .geometry import (STRAIGHT, Corridor, RailPlane, TrackFrame, fit_corridor, fit_rail_plane) from .fan_body import FanBody -from .lobula import Candidate, find_candidates, gauge_mask +from .lobula import Candidate, cluster_by_depth, find_candidates, gauge_mask from .medulla import EgoMotion, EgoMotionEstimator, EmdBank, looming from .mbon_readout import MbonReadout from .mushroom_body import Habituation, HabituationConfig, MushroomBody, NoMemory @@ -221,11 +221,12 @@ class Params: # сбрасывать никогда. reset_gap_s: float = 2.0 novelty_gate: float = 0.0 # ниже этой новизны кандидат отбрасывается сразу - # Где считать ламину: 'cpu', 'cuda' или 'auto' (видеокарта, если есть). - # По умолчанию процессор, и это решение, а не осторожность: GPU ускоряет - # только ламину (5.8 → 2.9 мс из 43), а на машине проверки требует - # другого контейнера и запуска с --gpus all. EXPERIMENTS п. 7.4. - device: str = "cpu" + # Где считать плотные стадии — сетчатку, ламину, кластеризацию (gpu.py): + # 'auto' — видеокарта, если PyTorch её видит, иначе процессор; 'cuda'; + # 'cpu'. Результат одинаков: на видеокарте те же операции в тех же типах. + # Любой сбой видеокарты переводит конвейер на процессор до конца работы, + # кадр досчитывается там же. EXPERIMENTS п. 21. + device: str = "auto" enable_motion: bool = True enable_memory: bool = True enable_looming: bool = False # T4/T5 + LPLC2: нужны для оценки надвигания и для @@ -271,6 +272,15 @@ class FlyGuard: self.p = params or Params() from .device import get_device self.device = get_device(self.p.device) + self.gpu = None # плотные стадии на видеокарте (gpu.GpuStages) + self.gpu_error: str | None = None + if self.device.startswith("cuda"): + try: + from .gpu import GpuStages + self.gpu = GpuStages(self.device) + self.gpu.warmup() + except Exception as e: # нет драйвера, не та видеокарта, нет памяти + self._gpu_off(e) self.memory = memory if (memory is not None and self.p.enable_memory) else NoMemory() self.readout = readout if (readout is not None and self.p.enable_mbon) else None self.layout_full = layout @@ -333,6 +343,25 @@ class FlyGuard: else DescendingNeurons(warn_evidence=0.0, clear_evidence=0.0, emergency_evidence=0.0, min_hits=1)) + # ------------------------------------------------------------------ видеокарта + + def _gpu_off(self, exc: Exception) -> None: + """Сбой видеокарты: до конца работы считать на процессоре.""" + from .device import notify_cuda_error + notify_cuda_error(exc) + self.gpu = None + self.device = "cpu" + self.gpu_error = f"{type(exc).__name__}: {exc}" + + def _cluster(self, mask: np.ndarray, r: np.ndarray, **kw): + """Кластеризация с учётом глубины: видеокарта, при сбое — процессор.""" + if self.gpu is not None: + try: + return self.gpu.cluster_by_depth(mask, r, **kw) + except Exception as e: + self._gpu_off(e) + return cluster_by_depth(mask, r, **kw) + # ------------------------------------------------------------------ калибровка def _ensure_layout(self, pc: PointCloud2) -> bool: @@ -412,7 +441,14 @@ class FlyGuard: self.frames_seen += 1 with t("retina"): - img: RangeImage = self.layout_full.project(pc, self.cols) + img: RangeImage | None = None + if self.gpu is not None: + try: + img = self.gpu.project(self.layout_full, pc, self.cols) + except Exception as e: + self._gpu_off(e) + if img is None: + img = self.layout_full.project(pc, self.cols) with t("stabilize"): self.plane = fit_rail_plane(img, self.layout, prev=self.plane) @@ -423,12 +459,14 @@ class FlyGuard: else STRAIGHT) with t("lamina"): - dev = self.device - if dev.startswith("cuda"): - from .device import is_cuda_available - if not is_cuda_available(): - self.device = dev = "cpu" - lam = lamina.process(tf.r, tf.valid, device=dev) + lam = None + if self.gpu is not None: + try: + lam = self.gpu.lamina() # по образу кадра, уже лежащему на видеокарте + except Exception as e: + self._gpu_off(e) + if lam is None: + lam = lamina.process(tf.r, tf.valid, device="cpu") with t("ego"): ego = (self.ego_est.update(tf, pc.stamp) if self.p.enable_motion @@ -475,7 +513,8 @@ class FlyGuard: h_top=self.p.h_top, half_width_top=self.p.half_width_top, top_d_max=self.p.top_d_max, - top_detect=self.p.top_detect) + top_detect=self.p.top_detect, + cluster=self._cluster) with t("mushroom"): cands = self.memory.annotate(cands) diff --git a/requirements.txt b/requirements.txt index 7968446..0a01f73 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,8 @@ scipy>=1.10 lightgbm>=4.0 # графики полигона: tools/plot_benchmark.py matplotlib>=3.7 -# только для обучения на GPU (tools/train_mbon.py --device cuda); ядру не нужен +# видеокарта: плотные стадии конвейера (flyguard/gpu.py) и обучение +# (tools/train_mbon.py --device cuda). Без него ядро считает на процессоре — +# с тем же результатом # torch>=2.0 pytest>=7.0 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 5f82dd1..30799ec 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -872,3 +872,110 @@ def test_gauge_outline_bends_with_the_track(): assert c100 == pytest.approx(-100.0 ** 2 / (2 * R), rel=1e-3) assert at100[:, 0].min() == pytest.approx(c100 - 1.2, abs=1e-4) assert at100[:, 0].max() == pytest.approx(c100 + 1.2, abs=1e-4) + + +# ------------------------------------------------------------ видеокарта + +def _cuda_or_skip(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("нет видеокарты с CUDA") + + +def _scene(seed=0, shape=(128, 600)): + """Образ, похожий на тоннель: гладкие стены, ступеньки-предметы, пропуски эха.""" + rng = np.random.default_rng(seed) + h, w = shape + r = (np.linspace(3.0, 180.0, w)[None, :] * (1.0 + 0.2 * np.sin(np.arange(h)[:, None] / 9.0)) + ).astype(np.float32) + for _ in range(40): + i, j = rng.integers(0, h - 8), rng.integers(0, w - 12) + r[i:i + rng.integers(2, 8), j:j + rng.integers(2, 12)] *= np.float32(rng.uniform(0.3, 0.9)) + r += rng.normal(0.0, 0.02, r.shape).astype(np.float32) + valid = rng.random(shape) > 0.1 + return np.where(valid, r, np.float32(0.0)), valid + + +def test_gpu_lamina_and_clustering_match_cpu_bit_for_bit(): + """Ламина и кластеризация на видеокарте — те же числа и те же номера компонент.""" + _cuda_or_skip() + import torch + from flyguard import lamina + from flyguard.gpu import GpuStages + from flyguard.lobula import cluster_by_depth + + g = GpuStages("cuda") + for seed in range(3): + r, valid = _scene(seed) + g.r = torch.as_tensor(r, device="cuda") + g.valid = torch.as_tensor(valid, device="cuda") + a, b = lamina._process_cpu(r, valid), g.lamina() + for f in ("disp", "on", "off", "on_scale", "surround", "hole"): + assert np.array_equal(getattr(a, f), getattr(b, f)), f + mask = valid & (np.random.default_rng(seed).random(r.shape) < 0.5) + la, na = cluster_by_depth(mask, r) + lb, nb = g.cluster_by_depth(mask, r) + assert na == nb and np.array_equal(la, lb) + + +def test_gpu_projection_matches_cpu(): + """Сетчатка на видеокарте: целый кадр — бит в бит, перемешанный — почти.""" + _cuda_or_skip() + from flyguard.bag import Bag + from flyguard.cdr import PointCloud2 + from flyguard.gpu import GpuStages + + if not DATA.exists(): + pytest.skip("нет записей") + bag = Bag(next(p for p in DATA.iterdir() if p.is_dir())) + clouds = [pc for _, pc in bag.frames(start=2, stop=16)] + lay = ScanLayout.calibrate(clouds[:12]) + cols = lay.column_slice(30.0) + g = GpuStages("cuda") + pc = clouds[-1] + a, b = lay.project(pc, cols), g.project(lay, pc, cols) + for f in ("r_near", "r_far", "inten", "valid"): + assert np.array_equal(getattr(a, f), getattr(b, f)), f + + # порядок точек сбит, как в синтетике организаторов: раскладка по углам + rng = np.random.default_rng(0) + mixed = PointCloud2(pc.stamp, pc.frame_id, 1, pc.n_points, pc.point_step, + pc.is_dense, pc.points[rng.permutation(pc.n_points)]) + a, b = lay.project(mixed, cols), g.project(lay, mixed, cols) + assert a.valid.sum() > 1000 + differ = sum(int((getattr(a, f) != getattr(b, f)).sum()) + for f in ("r_near", "r_far", "inten", "valid")) + assert differ < 1e-3 * a.valid.size # последний знак арктангенса и порядок равных + + +def test_pipeline_survives_gpu_failure_mid_frame(): + """Отказ видеокарты посреди кадра: кадр досчитан на процессоре, решение то же.""" + from flyguard.bag import Bag + from flyguard.device import reset_device_cache + from flyguard.pipeline import FlyGuard, Params + + if not DATA.exists(): + pytest.skip("нет записей") + + class Broken: + def project(self, *a, **k): + raise RuntimeError("CUDA error: device lost (имитация)") + + lamina = cluster_by_depth = project + + frames = [pc for _, pc in Bag(next(p for p in DATA.iterdir() if p.is_dir())).frames(stop=30)] + ref = FlyGuard(Params(device="cpu")) + fg = FlyGuard(Params(device="cpu")) + try: + for i, pc in enumerate(frames): + if i == 20: # видеокарта «была» и отказала + fg.gpu, fg.device = Broken(), "cuda" + a, b = ref.process(pc), fg.process(pc) + assert (a is None) == (b is None) + if a is not None: + assert a.decision.detected == b.decision.detected + assert [c.d for c in a.candidates] == [c.d for c in b.candidates] + assert fg.gpu is None and fg.device == "cpu" + assert "device lost" in fg.gpu_error + finally: + reset_device_cache()