Zhirik1337
7507a123e5
- flyguard/device.py: автоопределение NVIDIA GPU, сбор телеметрии и Graceful Fallback на CPU; - flyguard/lamina.py: ускорение 2D DoG фильтрации на тензорах PyTorch CUDA (0.25 мс вместо 8 мс); - flyguard/mbon_readout.py: векторизованный GPU-цикл обучения MBON без CPU-синхронизаций (поддержка 50k-100k клеток Кеньона); - flyguard/mushroom_body.py, flyguard/track_readout.py: аппаратное ускорение и опция GPU-бустинга; - flyguard/pipeline.py, tools/*.py: сквозная поддержка параметра device='auto' для всех инструментов; - Dockerfile.gpu, requirements-gpu.txt, docker-compose.yml: MLOps-инфраструктура под NVIDIA RTX 4070 Ti Super 16GB; - tests/test_pipeline.py: добавлены юнит-тесты на GPU device detection и роутинг ламины (40 тестов PASS).
380 lines
21 KiB
Python
380 lines
21 KiB
Python
"""MBON с обучением с учителем — считывание, а не новая сеть.
|
||
|
||
Грибовидное тело в `mushroom_body.py` учится **без меток**: синапсы KC→MBON
|
||
депрессируются на всём, что тоннель показывает часто, и выход MBON означает
|
||
«незнакомо». Это ровно familiarity suppression MBON-α′3 и ровно то, что нужно,
|
||
когда меток нет.
|
||
|
||
Но у мухи та же схема умеет и другое. При обучении с подкреплением
|
||
дофаминергические нейроны PPL1/PAM депрессируют KC→MBON **избирательно** — те
|
||
клетки, что были активны вместе с наказанием, — и выход MBON начинает означать
|
||
не «незнакомо», а «это предвещает удар». Один и тот же нейропиль, один и тот же
|
||
разрежённый код, другой учитель.
|
||
|
||
Здесь сделано именно это. Слои не меняются:
|
||
|
||
* вход — те же признаки кандидата плюс опора веерного тела (23 «проекционных
|
||
нейрона»);
|
||
* KC — та же случайная разрежённая проекция по 6 «когтей» на клетку;
|
||
* APL — то же глобальное торможение «победитель забирает всё», но отклик
|
||
остаётся **градуальным** и делится на общую активность (дивизивная
|
||
нормировка), а не превращается в единицы и нули;
|
||
* MBON — один выход, веса которого обучены различать предмет и обстановку.
|
||
|
||
Метки берутся не из разметки (её нет), а из физики: `flyguard.synth` вставляет
|
||
предмет трассировкой лучей, и кандидат считается предметом, если его ядро
|
||
состоит из лучей, в которые предмет действительно записан
|
||
(`tools/make_training_set.py`).
|
||
|
||
Зачем это нужно поверх ручной формулы веса улики. В `central_complex._quality`
|
||
шесть множителей, придуманных руками, а в дескрипторе 23 признака: **тень и
|
||
интенсивность в вес улики не входят вообще**, хотя окклюзионная тень за
|
||
предметом на большой дальности во много раз крупнее самого предмета.
|
||
|
||
Стоимость в инференсе — одно умножение матрицы 23 × n_kc на кандидата, доли
|
||
миллисекунды на CPU. GPU нужен только на обучении, и то не обязателен.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
from .lobula import Candidate
|
||
from .mushroom_body import describe
|
||
|
||
|
||
@dataclass
|
||
class MbonConfig:
|
||
"""Параметры считывания."""
|
||
|
||
# Ёмкость выбирается замером (`tools/train_mbon.py --sweep-kc`), а не на
|
||
# глаз: отбор «победитель забирает всё» стоит дороже матмула, и в худшем
|
||
# кадре с 64 кандидатами 20 000 клеток это 7.3 мс, 8 000 — 2.8 мс, 4 000 —
|
||
# 1.3 мс. Умолчание взято средним: платить временем кадра за ёмкость имеет
|
||
# смысл, только если развёртка показала, что она что-то даёт.
|
||
n_kc: int = 8_000 # клеток Кеньона в этом контуре
|
||
claws: int = 6 # входов на клетку (из коннектома)
|
||
sparsity: float = 0.0125 # доля активных после торможения APL — 100 клеток
|
||
seed: int = 20260921
|
||
|
||
|
||
def describe_full(c: Candidate) -> np.ndarray:
|
||
"""Вектор кандидата для считывания: дескриптор памяти плюс опора накопителя."""
|
||
acc = float(c.extra.get("acc_support", 0.0)) if c.extra else 0.0
|
||
return np.append(describe(c), np.float32(acc)).astype(np.float32)
|
||
|
||
|
||
class MbonReadout:
|
||
"""Разрежённый код клеток Кеньона + обученный линейный выход MBON."""
|
||
|
||
def __init__(self, cfg: MbonConfig | None = None, n_pn: int = 23):
|
||
self.cfg = cfg or MbonConfig()
|
||
self.n_pn = int(n_pn)
|
||
self.n_active = max(1, int(round(self.cfg.n_kc * self.cfg.sparsity)))
|
||
rng = np.random.default_rng(self.cfg.seed)
|
||
idx = np.stack([rng.choice(self.n_pn, size=self.cfg.claws, replace=False)
|
||
for _ in range(self.cfg.n_kc)])
|
||
sign = rng.choice((-1.0, 1.0), size=idx.shape)
|
||
w = np.zeros((self.cfg.n_kc, self.n_pn), np.float32)
|
||
np.put_along_axis(w, idx, sign.astype(np.float32), axis=1)
|
||
self.W = w
|
||
# Опорный отсчёт по полосам дальности: медиана калиброванного
|
||
# логита среди ВСЕХ кандидатов полосы. Нужен, чтобы отделить
|
||
# «насколько это предметно» от «как часто предметы вообще
|
||
# встречаются на такой дальности» — см. _prior_corr.
|
||
self.ref_edges = np.array([0., 30., 55., 80., 110., 160., 230.],
|
||
np.float32)
|
||
self.ref_z: np.ndarray | None = None
|
||
self.mean = np.zeros(self.n_pn, np.float32)
|
||
self.scale = np.ones(self.n_pn, np.float32)
|
||
self.w_mbon = np.zeros(self.cfg.n_kc, np.float32)
|
||
self.bias = np.float32(0.0)
|
||
# калибровка выхода: sigmoid(gain·z + shift). Градиентный спуск по
|
||
# разрежённому коду хорошо упорядочивает кандидатов, но масштаб логита
|
||
# зависит от ёмкости и числа эпох, а нам нужна осмысленная вероятность —
|
||
# она идёт множителем в вес улики.
|
||
self.gain = np.float32(1.0)
|
||
self.shift = np.float32(0.0)
|
||
|
||
# ------------------------------------------------------------------ код
|
||
|
||
def fit_normalizer(self, X: np.ndarray) -> None:
|
||
X = np.atleast_2d(np.asarray(X, np.float32))
|
||
self.mean = X.mean(0).astype(np.float32)
|
||
s = X.std(0).astype(np.float32)
|
||
self.scale = np.where(s > 1e-6, s, 1.0).astype(np.float32)
|
||
|
||
def encode(self, X: np.ndarray, device: str | None = None):
|
||
"""Признаки → (индексы активных клеток, их нормированный отклик).
|
||
|
||
Отклик остаётся градуальным и делится на свою сумму: это дивизивная
|
||
нормировка APL. Двоичный код тут заметно хуже — он выбрасывает
|
||
«насколько» клетка возбуждена, а для решения это как раз важно.
|
||
|
||
Нормировка приводит средний отклик к единице, а не сумму: иначе каждая
|
||
из сотни активных клеток даёт вклад 0.01, логиты выходят микроскопические
|
||
и обучение сваливается в разумное ранжирование при бессмысленной
|
||
калибровке (замерено: AUC 0.97 при доле верных 0.37).
|
||
"""
|
||
X = np.atleast_2d(np.asarray(X, np.float32))
|
||
if X.shape[1] != self.n_pn:
|
||
raise ValueError(f"считывание обучено на {self.n_pn} признаках, "
|
||
f"а дескриптор даёт {X.shape[1]}")
|
||
k = self.n_active
|
||
chunk = max(1, int(2 ** 26 // max(self.cfg.n_kc, 1)))
|
||
|
||
if device and device != "cpu":
|
||
import torch
|
||
out_i = np.empty((X.shape[0], k), np.int64)
|
||
out_v = np.empty((X.shape[0], k), np.float32)
|
||
with torch.no_grad():
|
||
m = torch.as_tensor(self.mean, device=device)
|
||
s = torch.as_tensor(self.scale, device=device)
|
||
w = torch.as_tensor(self.W, device=device).T.contiguous()
|
||
for i in range(0, X.shape[0], chunk):
|
||
t = torch.as_tensor(X[i:i + chunk], device=device)
|
||
y = torch.relu(((t - m) / s) @ w)
|
||
v, a = torch.topk(y, k, dim=1)
|
||
v = v * (k / v.sum(1, keepdim=True).clamp_min(1e-6))
|
||
out_i[i:i + chunk] = a.cpu().numpy()
|
||
out_v[i:i + chunk] = v.cpu().numpy()
|
||
return out_i, out_v
|
||
|
||
out_i = np.empty((X.shape[0], k), np.int64)
|
||
out_v = np.empty((X.shape[0], k), np.float32)
|
||
for i in range(0, X.shape[0], chunk):
|
||
z = (X[i:i + chunk] - self.mean) / self.scale
|
||
y = np.maximum(z @ self.W.T, 0.0)
|
||
a = np.argpartition(-y, k - 1, axis=1)[:, :k]
|
||
v = np.take_along_axis(y, a, axis=1)
|
||
v *= k / np.maximum(v.sum(1, keepdims=True), 1e-6)
|
||
out_i[i:i + chunk] = a
|
||
out_v[i:i + chunk] = v
|
||
return out_i, out_v
|
||
|
||
# ------------------------------------------------------------------ выход
|
||
|
||
def score(self, X: np.ndarray) -> np.ndarray:
|
||
"""Вероятность «это посторонний предмет», 0…1."""
|
||
a, v = self.encode(X)
|
||
z = self.logit(X, code=(a, v))
|
||
return 1.0 / (1.0 + np.exp(-z))
|
||
|
||
def logit(self, X: np.ndarray, code=None) -> np.ndarray:
|
||
a, v = code if code is not None else self.encode(X)
|
||
z = self.bias + (self.w_mbon[a] * v).sum(axis=1)
|
||
return self.gain * z + self.shift
|
||
|
||
def score_of(self, c: Candidate) -> float:
|
||
return float(self.score(describe_full(c)[None, :])[0])
|
||
|
||
def _ref_at(self, d: np.ndarray) -> np.ndarray:
|
||
"""Типичный отсчёт модели на дальности d (по обучающей выборке)."""
|
||
mid = 0.5 * (self.ref_edges[:-1] + self.ref_edges[1:])
|
||
return np.interp(d, mid, self.ref_z)
|
||
|
||
def _prior_corr(self, d: np.ndarray, prior_from: float) -> np.ndarray:
|
||
"""Поправка на распространённость предметов: логит → отношение.
|
||
|
||
Модель калибрована в апостериорную вероятность, а доля предметов
|
||
среди кандидатов падает с 42 % вблизи до 3.1 % за 160 м. Поэтому
|
||
честная вероятность на дальности низка даже у настоящего предмета:
|
||
замер даёт 0.051 у вставленного человека на 150 м — при 0.003 у
|
||
обстановки, то есть модель РАЗЛИЧАЕТ (AUC 0.743), но абсолютная
|
||
величина раздавлена редкостью. Детектору нужно отношение
|
||
правдоподобий, а не доля, поэтому из логита вычитается типичный
|
||
для этой полосы — ровно как число лучей уже нормируется на
|
||
ожидаемое для дальности. Ближе `prior_from` поправка нулевая.
|
||
"""
|
||
if self.ref_z is None or prior_from <= 0.0:
|
||
return np.zeros_like(d)
|
||
base = float(self._ref_at(np.array([prior_from], np.float32))[0])
|
||
corr = base - self._ref_at(d)
|
||
return np.where(d > prior_from, np.maximum(corr, 0.0), 0.0)
|
||
|
||
def annotate(self, cands: list[Candidate], prior_from: float = 0.0
|
||
) -> list[Candidate]:
|
||
if not cands:
|
||
return cands
|
||
X = np.stack([describe_full(c) for c in cands])
|
||
z = self.logit(X)
|
||
if prior_from > 0.0:
|
||
d = np.array([c.d for c in cands], np.float32)
|
||
z = z + self._prior_corr(d, prior_from)
|
||
for c, zz in zip(cands, z):
|
||
c.extra["mbon"] = float(1.0 / (1.0 + np.exp(-zz)))
|
||
return cands
|
||
|
||
# ------------------------------------------------------------------ обучение
|
||
|
||
def _learn_cpu(self, X: np.ndarray, y: np.ndarray, *, epochs: int = 60,
|
||
lr: float = 4.0, l2: float = 1e-5, verbose: bool = False) -> None:
|
||
"""Обучение на CPU через NumPy."""
|
||
y = np.asarray(y, np.float32)
|
||
a, v = self.encode(X, device="cpu")
|
||
n, k = a.shape
|
||
flat = a.ravel()
|
||
for ep in range(epochs):
|
||
z = self.bias + (self.w_mbon[a] * v).sum(axis=1)
|
||
p = 1.0 / (1.0 + np.exp(-z))
|
||
g = (p - y) / n
|
||
grad = np.bincount(flat, weights=np.repeat(g, k) * v.ravel(),
|
||
minlength=self.cfg.n_kc).astype(np.float32)
|
||
self.w_mbon -= lr * (grad + l2 * self.w_mbon)
|
||
self.bias -= np.float32(lr * g.sum())
|
||
if verbose and (ep + 1) % 50 == 0:
|
||
eps = 1e-7
|
||
loss = -(y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps)).mean()
|
||
print(f" эпоха {ep + 1:4d}: логистическая потеря {loss:.4f}")
|
||
zr = self.bias + (self.w_mbon[a] * v).sum(axis=1)
|
||
self._calibrate(zr, y)
|
||
zc = self.gain * zr + self.shift
|
||
dd = np.exp(np.asarray(X[:, 0], np.float64))
|
||
ref = []
|
||
for lo, hi in zip(self.ref_edges[:-1], self.ref_edges[1:]):
|
||
m = (dd >= lo) & (dd < hi)
|
||
ref.append(float(np.median(zc[m])) if m.sum() >= 50
|
||
else (ref[-1] if ref else 0.0))
|
||
self.ref_z = np.array(ref, np.float32)
|
||
|
||
def _learn_gpu(self, X: np.ndarray, y: np.ndarray, *, epochs: int = 60,
|
||
lr: float = 4.0, l2: float = 1e-5, device: str = "cuda",
|
||
verbose: bool = False) -> None:
|
||
"""GPU-ускоренная оптимизация MBON с учителем через тензоры PyTorch на CUDA.
|
||
|
||
Устраняет синхронизацию с CPU внутри цикла по эпохам, позволяя масштабировать
|
||
ёмкость до 50 000 – 100 000 клеток Кеньона без падения скорости.
|
||
"""
|
||
import torch
|
||
dev = torch.device(device)
|
||
y_t = torch.as_tensor(y, dtype=torch.float32, device=dev)
|
||
n = y_t.shape[0]
|
||
k = self.n_active
|
||
|
||
m = torch.as_tensor(self.mean, dtype=torch.float32, device=dev)
|
||
s = torch.as_tensor(self.scale, dtype=torch.float32, device=dev)
|
||
w_proj = torch.as_tensor(self.W, dtype=torch.float32, device=dev).T.contiguous()
|
||
|
||
chunk = max(1, int(2 ** 26 // max(self.cfg.n_kc, 1)))
|
||
a_list, v_list = [], []
|
||
with torch.no_grad():
|
||
for i in range(0, n, chunk):
|
||
t = torch.as_tensor(X[i:i + chunk], dtype=torch.float32, device=dev)
|
||
y_proj = torch.relu(((t - m) / s) @ w_proj)
|
||
v_chunk, a_chunk = torch.topk(y_proj, k, dim=1)
|
||
v_chunk = v_chunk * (k / v_chunk.sum(1, keepdim=True).clamp_min(1e-6))
|
||
a_list.append(a_chunk)
|
||
v_list.append(v_chunk)
|
||
|
||
a_t = torch.cat(a_list, dim=0) if len(a_list) > 1 else a_list[0]
|
||
v_t = torch.cat(v_list, dim=0) if len(v_list) > 1 else v_list[0]
|
||
|
||
w_mbon_t = torch.as_tensor(self.w_mbon, dtype=torch.float32, device=dev)
|
||
bias_t = torch.tensor(self.bias, dtype=torch.float32, device=dev)
|
||
flat_a = a_t.reshape(-1)
|
||
|
||
for ep in range(epochs):
|
||
z = bias_t + torch.sum(w_mbon_t[a_t] * v_t, dim=1)
|
||
p = torch.sigmoid(z)
|
||
g = (p - y_t) / n
|
||
weights = (g.unsqueeze(1) * v_t).reshape(-1)
|
||
grad = torch.zeros(self.cfg.n_kc, dtype=torch.float32, device=dev)
|
||
grad.scatter_add_(0, flat_a, weights)
|
||
w_mbon_t -= lr * (grad + l2 * w_mbon_t)
|
||
bias_t -= lr * g.sum()
|
||
|
||
if verbose and (ep + 1) % 50 == 0:
|
||
loss = -(y_t * torch.log(p.clamp_min(1e-7)) + (1.0 - y_t) * torch.log((1.0 - p).clamp_min(1e-7))).mean()
|
||
print(f" [GPU] эпоха {ep + 1:4d}: логистическая потеря {loss.item():.4f}")
|
||
|
||
zr_t = bias_t + torch.sum(w_mbon_t[a_t] * v_t, dim=1)
|
||
|
||
# Калибровка Платта полностью на GPU
|
||
g_t = torch.tensor(1.0, dtype=torch.float32, device=dev)
|
||
sh_t = torch.tensor(0.0, dtype=torch.float32, device=dev)
|
||
for _ in range(400):
|
||
p_cal = torch.sigmoid(g_t * zr_t + sh_t)
|
||
e_cal = p_cal - y_t
|
||
g_t -= 2.0 * (e_cal * zr_t).mean() / torch.clamp_min((zr_t * zr_t).mean(), 1e-6)
|
||
sh_t -= 2.0 * e_cal.mean()
|
||
|
||
self.w_mbon = w_mbon_t.cpu().numpy()
|
||
self.bias = np.float32(bias_t.item())
|
||
self.gain = np.float32(g_t.item())
|
||
self.shift = np.float32(sh_t.item())
|
||
|
||
zr = zr_t.cpu().numpy()
|
||
zc = self.gain * zr + self.shift
|
||
dd = np.exp(np.asarray(X[:, 0], np.float64))
|
||
ref = []
|
||
for lo, hi in zip(self.ref_edges[:-1], self.ref_edges[1:]):
|
||
m_mask = (dd >= lo) & (dd < hi)
|
||
ref.append(float(np.median(zc[m_mask])) if m_mask.sum() >= 50
|
||
else (ref[-1] if ref else 0.0))
|
||
self.ref_z = np.array(ref, np.float32)
|
||
|
||
def learn(self, X: np.ndarray, y: np.ndarray, *, epochs: int = 60,
|
||
lr: float = 4.0, l2: float = 1e-5, device: str | None = None,
|
||
verbose: bool = False) -> None:
|
||
"""Логистическая регрессия по разрежённому коду — депрессия с учителем.
|
||
|
||
Поддерживает прозрачное переключение между NVIDIA GPU (CUDA) и CPU.
|
||
"""
|
||
target_dev = device
|
||
if target_dev is None or target_dev == "auto":
|
||
from .device import get_device
|
||
target_dev = get_device("auto")
|
||
|
||
if target_dev.startswith("cuda"):
|
||
try:
|
||
self._learn_gpu(X, y, epochs=epochs, lr=lr, l2=l2,
|
||
device=target_dev, verbose=verbose)
|
||
return
|
||
except Exception as e:
|
||
import logging
|
||
logging.getLogger("flyguard.mbon").warning("GPU learning failed (%s), fallback to CPU", e)
|
||
|
||
self._learn_cpu(X, y, epochs=epochs, lr=lr, l2=l2, verbose=verbose)
|
||
|
||
def _calibrate(self, z: np.ndarray, y: np.ndarray, iters: int = 400) -> None:
|
||
"""Шкалирование Платта: подобрать наклон и сдвиг по обучающей выборке."""
|
||
g, sh = 1.0, 0.0
|
||
for _ in range(iters):
|
||
p = 1.0 / (1.0 + np.exp(-(g * z + sh)))
|
||
e = p - y
|
||
g -= 2.0 * float((e * z).mean()) / max(float((z * z).mean()), 1e-6)
|
||
sh -= 2.0 * float(e.mean())
|
||
self.gain, self.shift = np.float32(g), np.float32(sh)
|
||
|
||
# ------------------------------------------------------------------ хранение
|
||
|
||
def save(self, path: str | Path) -> None:
|
||
np.savez_compressed(path, w_mbon=self.w_mbon, bias=self.bias,
|
||
gain=self.gain, shift=self.shift,
|
||
mean=self.mean, scale=self.scale,
|
||
n_kc=self.cfg.n_kc, claws=self.cfg.claws,
|
||
sparsity=self.cfg.sparsity, seed=self.cfg.seed,
|
||
n_pn=self.n_pn, ref_edges=self.ref_edges,
|
||
ref_z=(self.ref_z if self.ref_z is not None
|
||
else np.zeros(0, np.float32)))
|
||
|
||
@classmethod
|
||
def load(cls, path: str | Path) -> "MbonReadout":
|
||
"""Загрузка через `cls`, а не через имя класса: считывание по трекам —
|
||
тот же контур с другим входом, и оно наследуется отсюда целиком."""
|
||
d = np.load(path, allow_pickle=False)
|
||
cfg = MbonConfig(n_kc=int(d["n_kc"]), claws=int(d["claws"]),
|
||
sparsity=float(d["sparsity"]), seed=int(d["seed"]))
|
||
m = cls(cfg, n_pn=int(d["n_pn"]))
|
||
m.w_mbon = d["w_mbon"].astype(np.float32)
|
||
m.bias = np.float32(d["bias"])
|
||
m.gain = np.float32(d["gain"])
|
||
m.shift = np.float32(d["shift"])
|
||
m.mean = d["mean"].astype(np.float32)
|
||
m.scale = d["scale"].astype(np.float32)
|
||
if "ref_z" in d.files and d["ref_z"].size:
|
||
m.ref_edges = d["ref_edges"].astype(np.float32)
|
||
m.ref_z = d["ref_z"].astype(np.float32)
|
||
return m
|