mllightitup/Stgnf.md

## Stgnf.md

      
    Raw
  

              Stgnf.md
            
          
    STG-NF: Полная Архитектура и Спецификация


Spatio-Temporal Graph Normalizing Flows для обнаружения аномалий в позах человека.
Конфигурация: ShanghaiTech, uniform strategy, стандартные параметры.


1. Сводка Параметров (args.py defaults)


Параметр
Значение
Описание


dataset
ShanghaiTech
Датасет


seg_len (τ)
24
Длина сегмента позы (кадры)


seg_stride
6
Шаг скользящего окна (train). Test = 1


K
8
Количество Flow Steps


L
1
Количество уровней (без Squeeze)


R
3.0
Среднее prior-распределения μ_normal


hidden_channels
0
Скрытые каналы ST-GCN (0 = один блок)


flow_permutation
permute
Тип перестановки (reverse order)


flow_coupling
affine
Тип coupling layer


adj_strategy
uniform
Стратегия графа


max_hops
8
Макс. дистанция в adjacency


edge_importance
False
Веса рёбер


temporal_kernel
None → 13
Авто: T//2 + 1 = 24//2+1 = 13


model_confidence
False
Не использовать канал confidence


headless
False
Все 18 keypoints


batch_size
256
Размер батча


epochs
8
Число эпох


model_lr
5e-4
Learning rate


model_optimizer
adamx
Optimax optimizer


model_sched
exp_decay → None
Scheduler = None (ручной decay)


model_lr_decay
0.99
LR decay per epoch


model_weight_decay
5e-5
Weight decay


LU_decomposed
True
LU-разложение для InvConv


learn_top
False
Без обучаемого prior


actnorm_scale
1.0
Масштаб ActNorm


num_transform
2
Число аугментаций (identity + flip)


prop_norm_scale
1
Пропорциональная нормализация


norm_scale
0
Не масштабировать непропорционально


global_pose_segs
True
Нормализовать позы


seed
999
Random seed


2. Общая Блок-Схема Выполнения


      graph TD
    A["main() — train_eval.py"] --> B["Парсинг аргументов<br/>init_parser() + init_sub_args()"]
    B --> C["Установка seed"]
    C --> D["Создание директории эксперимента<br>create_exp_dirs()"]
    D --> E["Загрузка данных<br>get_dataset_and_loader()"]
    E --> F["Инициализация параметров модели<br>init_model_params()"]
    F --> G["Создание модели STG_NF"]
    G --> H["Создание Trainer"]
    H --> I{"pretrained?"}
    I -- Нет --> J["trainer.train()"]
    I -- Да --> K["trainer.load_checkpoint()"]
    J --> L["trainer.test()"]
    K --> L
    L --> M["score_dataset() — подсчёт AuC"]
    M --> N["Вывод результата"]

    
      Loading

  
3. Data Pipeline

3.1 Структура Входных Данных

Данные загружаются из JSON-файлов формата AlphaPose tracked person:
data/ShanghaiTech/pose/train/  — JSON файлы
data/ShanghaiTech/pose/test/   — JSON файлы

Каждый JSON имеет формат:
{
  "person_id": {
    "frame_id": {
      "keypoints": [x1, y1, c1, x2, y2, c2, ..., x17, y17, c17],  // 51 float
      "scores": float
    }
  }
}
3.2 Блок-Схема Data Pipeline


      graph TD
    A["JSON файлы<br/>(AlphaPose tracked persons)"] --> B["gen_dataset()"]
    B --> C["Для каждого JSON:<br/>gen_clip_seg_data_np()"]
    C --> D["single_pose_dict2np()<br/>Извлечение поз каждого человека"]
    D --> E["split_pose_to_segments()<br/>Скользящее окно: seg_len=24, stride=6"]
    E --> F["Сегменты shape: (N_seg, 24, 17, 3)<br/>x, y, confidence"]
    F --> G{"17 кейпоинтов?"}
    G -- Да --> H["keypoints17_to_coco18()<br/>17→18 кейпоинтов"]
    H --> I["Транспонирование<br/>(N, 3, 24, 18) = (N, C, T, V)<br/>dtype: float32"]
    I --> J["PoseSegDataset"]
    J --> K["__getitem__:<br/>Выбрать 2 канала [:, :2]<br/>= (2, 24, 18)"]
    K --> L["normalize_pose()<br/>Нормализация"]
    L --> M["Готовый тензор<br/>shape: (2, 24, 18)<br/>dtype: float32"]

    
      Loading

  
3.3 Преобразование 17→18 Keypoints (keypoints17_to_coco18)

# Добавляем neck = среднее плеч (kp[5] + kp[6]) / 2
neck = 0.5 * (kp[..., 5, :] + kp[..., 6, :])
kp = concat([kp, neck], axis=-2)  # 18 kp
# Переставляем порядок:
order = [0, 17, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3]
kp_coco18 = kp[..., order, :]
Результат: OpenPose 18-keypoint формат: [Nose, Neck, RShoulder, RElbow, RWrist, LShoulder, LElbow, LWrist, RHip, RKnee, RAnkle, LHip, LKnee, LAnkle, REye, LEye, REar, LEar]
3.4 Нормализация Позы (normalize_pose)

Вход: (N, T, V, F) = (N, 24, 18, 3) где F = [x, y, confidence]
def normalize_pose(pose_data):
    vid_res = [856, 480]           # разрешение видео ShanghaiTech
    norm_factor = [856, 480, 1]    # для [x, y, conf]
    
    # 1. Нормализация по разрешению видео
    pose = pose_data / norm_factor  # x/856, y/480, conf/1
    
    # 2. Центрирование + масштабирование по std
    #    mean по осям (T, V) для каждого сэмпла
    mean_xy = pose[..., :2].mean(axis=(1, 2))       # (N, 2)
    std_y = pose[..., 1].std(axis=(1, 2))            # (N,)
    
    pose[..., :2] = (pose[..., :2] - mean_xy[:, None, None, :]) / std_y[:, None, None, None]
    # confidence канал остаётся нормализованным к [0, 1]
    
    return pose  # (N, 24, 18, 3)
Important
При model_confidence=False (по умолчанию) перед подачей в модель берутся только первые 2 канала: samp = data[0][:, :2] → shape (B, 2, 24, 18).
Канал confidence отбрасывается. Это критично для ONNX-экспорта.

3.5 Аугментация (train only)

Используются num_transform=2 аугментации из trans_list:

Transform 0: Identity (без изменений)
Transform 1: Horizontal flip (отражение по X)

Каждая аугментация — аффинная трансформация координат (x, y):
# flip: trans_mat = [[-1, 0, 0], [0, 1, 0], [0, 0, 1]]
# Применяется к [x, y, 1], confidence сохраняется
# Dataset.__len__ = num_transform * num_samples = 2 * N
3.6 Финальная Форма Тензоров


Стадия
Shape
Dtype
Описание


Raw JSON keypoints
(51,) flat
float
[x1,y1,c1,...,x17,y17,c17]


Single person track
(T_clip, 17, 3)
float64
Все кадры 1 человека


After segmentation
(N_seg, 24, 17, 3)
float64
Скользящее окно


After 17→18
(N_seg, 24, 18, 3)
float64
+neck keypoint


After transpose
(N_seg, 3, 24, 18) = (N, C, T, V)
float32
Каналы первые


After [:, :2]
(B, 2, 24, 18)
float32
Без confidence


After normalize
(B, 2, 24, 18)
float32
Вход модели


4. Архитектура Графа (Adjacency Matrix)

4.1 Построение Графа для uniform + max_hops=8


      graph LR
    A["OpenPose 18 joints"] --> B["Self-links: (i,i) × 18"]
    A --> C["Anatomical edges: 17 пар"]
    B --> D["get_hop_distance()<br/>max_hop=8"]
    C --> D
    D --> E["adjacency matrix 18×18<br/>все пары с distance ≤ 8"]
    E --> F["normalize_digraph()<br/>AD = A · D^-1"]
    F --> G["A shape: (1, 18, 18)<br/>uniform: 1 матрица"]

    
      Loading

  
4.2 Детали Adjacency Matrix

При max_hops=8 и 18 joints (макс. расстояние в скелете ~7 рёбер) — все joints связаны друг с другом. Adjacency matrix — полносвязная (с self-loops).
# Нормализация: AD = A · D^(-1)
# D — диагональная матрица степеней (число входящих рёбер)
# Каждый столбец нормализован: сумма = 1
# → каждый joint влияет на соседей с равным весом
Shape A: (1, 18, 18) — float32, на GPU. K_spatial = A.size(0) = 1.
Tip
При uniform + max_hops=8 матрица смежности фактически эквивалентна Global Average Pooling по пространственному измерению. Каждый столбец A нормализован так, что операция x @ A[0] — это взвешенное среднее по всем joints.


5. Архитектура Модели STG-NF

5.1 Иерархия Классов


      classDiagram
    class STG_NF {
        +FlowNet flow
        +float R = 3.0
        +bool learn_top = False
        +Tensor prior_h (1, 4, 24, 18) zeros
        +Tensor prior_h_normal (4, 24, 18)
        +normal_flow(x, label, score) → z, nll
        +prior(data, label) → mean, logs
    }
    
    class FlowNet {
        +ModuleList layers
        +int K = 8
        +encode(z, logdet) → z, logdet
    }
    
    class FlowStep {
        +ActNorm2d actnorm
        +Permute2d reverse (permutation)
        +ModuleList block (ST-GCN)
        +Tensor A (1, 18, 18)
        +normal_flow(input, logdet) → z, logdet
    }
    
    class st_gcn {
        +ConvTemporalGraphical gcn
        +Sequential tcn
        +residual
        +ReLU relu
    }
    
    class ConvTemporalGraphical {
        +Conv2d conv
        +int kernel_size
        +forward(x, A) → x, A
    }
    
    STG_NF --> FlowNet
    FlowNet --> FlowStep : "×8 (K=8)"
    FlowStep --> ActNorm2d
    FlowStep --> Permute2d
    FlowStep --> st_gcn : "×1 (hidden=0)"
    st_gcn --> ConvTemporalGraphical

    
      Loading

  
5.2 Полная Блок-Схема Модели (Forward Pass)


      graph TD
    INPUT["Input x<br/>(B, 2, 24, 18)<br/>float32"] --> FLOW["FlowNet.encode()"]
    
    subgraph FlowNet ["FlowNet — K=8 FlowSteps"]
        direction TB
        FS1["FlowStep 0 (first=True)"]
        FS2["FlowStep 1"]
        FS3["FlowStep 2..6"]
        FS4["FlowStep 7 (last=True)"]
        FS1 --> FS2 --> FS3 --> FS4
    end
    
    FLOW --> FS1
    FS4 --> Z["z output<br/>(B, 2, 24, 18)"]
    Z --> PRIOR["prior() → mean, logs"]
    PRIOR --> GAUSS["gaussian_likelihood(mean, logs, z)"]
    GAUSS --> NLL["nll = (-objective) / (log2 * C * T * V)<br/>= (-objective) / (log2 * 2 * 24 * 18)"]

    
      Loading

  
5.3 Детальная Блок-Схема FlowStep (Affine Coupling)

Каждый из 8 FlowStep имеет одинаковую структуру:

  
      graph TD
    IN["Input z<br/>(B, 2, 24, 18)"] --> AN["1. ActNorm2d<br/>z = z * exp(logs) + bias<br/>logdet += sum(logs) * T * V"]
    AN --> PERM["2. Permute2d (reverse)<br/>Каналы: [1, 0] → [0 ↔ 1]<br/>logdet не меняется"]
    PERM --> SPLIT["3. split_feature(z, 'split')<br/>z1 = z[:, :1] → (B, 24, 18)<br/>z2 = z[:, 1:] → (B, 24, 18)"]
    SPLIT --> UNSQ["z1.unsqueeze(1) → (B, 1, 24, 18)<br/>z2.unsqueeze(1) → (B, 1, 24, 18)"]
    UNSQ --> CLONE["h = z1.clone()"]
    CLONE --> STGCN["ST-GCN block(h, A)<br/>in: (B, 1, 24, 18)<br/>out: (B, 2, 24, 18)"]
    STGCN --> CROSS["split_feature(h, 'cross')<br/>shift = h[:, 0::2] → (B, 24, 18)<br/>scale = h[:, 1::2] → (B, 24, 18)"]
    CROSS --> SCALE_FN["scale = sigmoid(scale + 2.0) + 1e-6<br/>Диапазон: (1e-6, 1+1e-6)"]
    SCALE_FN --> AFF["z2 = (z2 + shift) * scale"]
    AFF --> LOGDET["logdet += sum(log(scale), dims=[1,2,3])"]
    LOGDET --> CAT["z = cat(z1, z2, dim=1)<br/>→ (B, 2, 24, 18)"]
    CAT --> OUT["Output z<br/>(B, 2, 24, 18)"]

    
      Loading

  
Important
Split Feature с squeeze: Функция split_feature при type="split" делит тензор по dim=1 и вызывает .squeeze(dim=1). Для C=2: z1 = z[:, :1, :, :].squeeze(1) → (B, 24, 18) (3D). Потом перед ST-GCN делается unsqueeze(dim=1) → (B, 1, 24, 18).
При type="cross": shift = h[:, 0::2, :, :].squeeze(1), scale = h[:, 1::2, :, :].squeeze(1).


6. Детальная Архитектура Каждого Компонента

6.1 ActNorm2d

Activation Normalization — обучаемые bias и scale per channel. Инициализируется data-dependent на первом батче.
Параметры:
  bias:  (1, C, 1, 1) = (1, 2, 1, 1)  — init: -mean(first_batch)
  logs:  (1, C, 1, 1) = (1, 2, 1, 1)  — init: log(scale / std(first_batch))

Forward (not reverse):
  z = input + bias                     # center
  z = z * exp(logs)                    # scale  
  logdet += sum(logs) * T * V          # = sum(logs) * 24 * 18 = sum(logs) * 432

Число параметров: 2 * C = 2 * 2 = 4 (bias + logs)

6.2 Permute2d (reverse order permutation)

При flow_permutation='permute' — простая перестановка каналов в обратном порядке.
indices = [1, 0]       # для C=2
indices_inverse = [1, 0]

Forward:
  z = z[:, [1, 0], :, :]    # свап каналов x ↔ y
  logdet не меняется (перестановка = ортогональная матрица, det=±1, log|det|=0)

Note
Для C=2 reverse permute = swap channels. Это не InvertibleConv1x1. Нет обучаемых параметров. LogDet = 0.

6.3 Affine Coupling — ST-GCN Block

При hidden_channels=0 создаётся один блок ST-GCN:
# get_stgcn(in_channels=1, hidden_channels=0, out_channels=2,
#           temporal_kernel_size=13, spatial_kernel_size=1, first=True/False)

# hidden_channels == 0 → один блок:
block = ModuleList([
    st_gcn(in_channels=1, out_channels=2, kernel_size=(13, 1), stride=1, residual=False/True)
])
Первый FlowStep (k=0): first=True → residual=False → self.residual = lambda x: 0
Последующие FlowSteps (k=1..7): first=False → residual=True. Т.к. in_channels(1) != out_channels(2) → residual = Conv2d(1, 2, 1×1) + BN2d(2).
6.4 ST-GCN Block (st_gcn)


      graph TD
    X["Input x<br/>(B, 1, 24, 18)"] --> RES["residual(x)"]
    X --> GCN["ConvTemporalGraphical<br/>gcn(x, A)"]
    GCN --> TCN["TCN Sequential:<br/>BatchNorm2d(2)<br/>ReLU<br/>Conv2d(2→2, kernel=(13,1), pad=(6,0))<br/>BatchNorm2d(2)"]
    TCN --> ADD["x = tcn_out + residual"]
    RES --> ADD
    ADD --> RELU["ReLU"]
    RELU --> OUT["Output<br/>(B, 2, 24, 18)"]

    
      Loading

  
6.4.1 ConvTemporalGraphical (Spatial GCN)

kernel_size = 1 (spatial, от A.size(0)=1 для uniform)

Conv2d: in_channels=1, out_channels=2*1=2, kernel=(1,1), bias=True
  → (B, 1, 24, 18) → (B, 2, 24, 18)

# Graph convolution via einsum:
x = conv(x)                               # (B, 2, 24, 18)
n, kc, t, v = x.size()                    # B, 2, 24, 18
x = x.view(n, 1, 2, t, v)                 # (B, K=1, C=2, 24, 18)
x = einsum('nkctv,kvw->nctw', x, A)       # A: (1, 18, 18)
# Результат: (B, 2, 24, 18)

Tip
Для ONNX/torch.compile: torch.einsum('nkctv,kvw->nctw', x, A) при K=1 эквивалентен:
x = x.squeeze(1)  # (B, 2, 24, 18) затем x = torch.matmul(x, A[0])  # (B, 2, 24, 18) @ (18, 18).
Это можно заменить на F.linear или обычный torch.mm broadcasting.

6.4.2 Temporal Convolution Network (TCN)

Sequential:
  1. BatchNorm2d(2)        — нормализация по 2 каналам
  2. ReLU(inplace=True)
  3. Conv2d(2, 2, kernel_size=(13, 1), stride=(1, 1), padding=(6, 0))
  4. BatchNorm2d(2)

Input/Output: (B, 2, 24, 18) → (B, 2, 24, 18)
Padding = (13-1)//2 = 6 — same-padding по времени

6.4.3 Residual Connection


FlowStep k
residual


k=0 (first)
lambda x: 0 — нет residual


k=1..7
Conv2d(1→2, 1×1) + BN2d(2) — проекция каналов


6.5 Prior Distribution

# learn_top = False → prior фиксирован
# R = 3.0

# prior_h: buffer (1, 4, 24, 18) — zeros
# prior_h_normal: buffer (4, 24, 18):
#   [0:2] = ones * 3.0     # mean = R = 3.0
#   [2:4] = zeros           # logs = 0 → variance = exp(0*2) = 1

# prior() splits:
#   mean = h[:, :2, :, :] = R = 3.0  (unsupervised)
#   logs = h[:, 2:, :, :] = 0.0

# Для unsupervised (label=None):
#   h = prior_h.repeat(B, 1, 1, 1) = zeros(B, 4, 24, 18)
#   split → mean=(B, 2, 24, 18) zeros, logs=(B, 2, 24, 18) zeros
Warning
Нюанс: при label=None (unsupervised inference без лейблов) используется prior_h (zeros), а не prior_h_normal (R=3.0). Однако при обучении вызывается self.model(samp, label=label, score=score) где label = data[-1] — это ones для нормальных данных. В prior():
if label is not None:
    h[label == 1] = self.prior_h_normal  # mean=3.0, logs=0
При тестировании: label=torch.ones(B) → mean=3, logs=0.

6.6 Gaussian Likelihood

def gaussian_p(mean, logs, x):
    c = log(2π)
    return -0.5 * (logs * 2.0 + ((x - mean)**2) / (exp(logs * 2.0) + 1e-6) + c)

def gaussian_likelihood(mean, logs, x):
    # logs=0 → variance=1
    # = -0.5 * (0 + (z - mean)^2 / (1 + 1e-6) + log(2π))
    # ≈ -0.5 * ((z - 3.0)^2 + 1.8379)
    return sum(gaussian_p(mean, logs, x), dim=[1, 2, 3])
    # → скаляр на каждый элемент батча: (B,)
6.7 NLL Computation

def normal_flow(self, x, label, score):
    b, c, t, v = x.shape   # B, 2, 24, 18
    
    z, objective = self.flow(x, reverse=False)
    # objective = logdet из FlowNet (сумма всех logdet от 8 FlowSteps)
    
    mean, logs = self.prior(x, label)    # mean=3.0, logs=0.0
    objective += gaussian_likelihood(mean, logs, z)
    # objective = logdet + log p(z | μ=3, σ=1)
    
    nll = (-objective) / (log(2.0) * c * t * v)
    #    = (-objective) / (0.6931 * 2 * 24 * 18)
    #    = (-objective) / 598.45
    # → bits per dimension
    
    return z, nll   # z: (B, 2, 24, 18), nll: (B,)

7. Полная Таблица Параметров Модели

7.1 FlowStep k=0 (first, no residual)


Компонент
Параметры
Shape
Count


ActNorm


bias
learnable
(1, 2, 1, 1)
2


logs
learnable
(1, 2, 1, 1)
2


Permute2d
—
indices only
0


ST-GCN


gcn.conv.weight
learnable
(2, 1, 1, 1)
2


gcn.conv.bias
learnable
(2,)
2


tcn.BN1.weight
learnable
(2,)
2


tcn.BN1.bias
learnable
(2,)
2


tcn.Conv2d.weight
learnable
(2, 2, 13, 1)
52


tcn.Conv2d.bias
learnable
(2,)
2


tcn.BN2.weight
learnable
(2,)
2


tcn.BN2.bias
learnable
(2,)
2


Residual
None (=0)

0


Subtotal


68


7.2 FlowStep k=1..7 (with residual projection)


Компонент
Параметры
Shape
Count


ActNorm
bias + logs
2 × (1,2,1,1)
4


Permute2d
—

0


ST-GCN (same as k=0)


64


Residual Conv2d
weight
(2, 1, 1, 1)
2


Residual Conv2d
bias
—
(no bias, BN follows)


Residual BN2d
weight+bias
(2,) × 2
4


Subtotal


74


7.3 Общее Количество Параметров

FlowStep 0:  68 параметров
FlowStep 1-7: 7 × 74 = 518 параметров
─────────────────────────────
TOTAL: 68 + 518 = 586 обучаемых параметров

+ Buffers (не обучаемые):
  prior_h:        (1, 4, 24, 18) = 1728
  prior_h_normal: (4, 24, 18) = 1728  
  prior_h_abnormal: (4, 24, 18) = 1728
  A per FlowStep: 8 × (1, 18, 18) = 2592
  BN running_mean/var: 8 × (2 + 2 + 2 + 2) × 2 = varies

Note
Авторы заявляют ~1K параметров. Реальное число: ~586 обучаемых параметров + BatchNorm running statistics.


8. Training Pipeline

8.1 Блок-Схема Обучения


      graph TD
    A["model.train()<br/>model.to('cuda:0')"] --> B["Epoch loop: 0..7"]
    B --> C["DataLoader: batch_size=256, shuffle=True"]
    C --> D["Для каждого батча:"]
    D --> E["data → cuda, non_blocking"]
    E --> F["score = data[-2].amin(dim=-1)<br/>label = data[-1]<br/>samp = data[0][:, :2]"]
    F --> G["z, nll = model(samp.float(),<br/>label=label, score=score)"]
    G --> H["losses = mean(nll)"]
    H --> I["losses.backward()"]
    I --> J["clip_grad_norm_(params, max=100)"]
    J --> K["optimizer.step()"]
    K --> L["optimizer.zero_grad()"]
    L --> C
    B --> M["save_checkpoint()"]
    M --> N["adjust_lr: lr *= 0.99"]

    
      Loading

  
8.2 Optimizer & Scheduler

# Optimizer: Adamax
optimizer = optim.Adamax(model.parameters(), lr=5e-4)
# weight_decay передаётся через init_optimizer → partial, 
# но init_optimizer вызывается с lr=5e-4 only

# Scheduler: exp_decay → init_scheduler returns None
# → adjust_lr() использует ручной decay:
# new_lr = lr * (0.99 ** epoch)
# Epoch 0: lr = 5e-4
# Epoch 1: lr = 4.95e-4
# ...
# Epoch 7: lr = 5e-4 * 0.99^7 ≈ 4.66e-4
8.3 ActNorm Data-Dependent Initialization

При первом forward pass (первый батч первой эпохи) каждый ActNorm2d инициализируется:
def initialize_parameters(self, input):
    # input: (B, 2, 24, 18)
    bias = -mean(input, dim=[0, 2, 3], keepdim=True)  # (1, 2, 1, 1)
    vars = mean((input + bias)^2, dim=[0, 2, 3], keepdim=True)
    logs = log(1.0 / (sqrt(vars) + 1e-6))              # (1, 2, 1, 1)
Warning
ONNX-совместимость: ActNorm должен быть инициализирован ДО экспорта. При загрузке чекпоинта вызывается model.set_actnorm_init() который ставит m.inited = True. Для ONNX-экспорта нужно гарантировать что inited=True у всех ActNorm, иначе будет ошибка в eval mode.


9. Inference Pipeline

9.1 Блок-Схема Инференса


      graph TD
    A["model.eval()<br/>model.to('cuda:0')"] --> B["TestDataLoader:<br/>batch_size=256, stride=1"]
    B --> C["Для каждого батча:"]
    C --> D["samp = data[0][:, :2]<br/>float32, to cuda"]
    D --> E["torch.no_grad()"]
    E --> F["z, nll = model(samp,<br/>label=ones(B), score=score)"]
    F --> G["probs += (-1 * nll)"]
    G --> C
    B --> H["prob_mat_np = probs.cpu().numpy()"]
    H --> I["score_dataset()"]

    
      Loading

  
9.2 Score Dataset (Подсчёт AuC)


      graph TD
    A["Normality scores<br/>per segment (N_test,)"] --> B["get_dataset_scores()"]
    B --> C["Для каждого клипа:"]
    C --> D["Загрузка GT mask: .npy<br/>(1=normal, 0=abnormal)"]
    D --> E["Для каждого person_id:"]
    E --> F["Присвоение score кадрам:<br/>scores[frame + seg_len/2] = pid_score<br/>offset = seg_len/2 = 12"]
    F --> G["clip_score = min over all persons<br/>per frame"]
    G --> H["Replace inf with max/min"]
    H --> I["smooth_scores()<br/>Gaussian filter σ=1..6"]
    I --> J["roc_auc_score(gt, scores)"]
    J --> K["AuC result"]

    
      Loading

  
Important
Frame scoring: Каждому сегменту [t, t+24) присваивается 1 score. Этот score приписывается к кадру t + 12 (середина сегмента). Для кадров с несколькими людьми берётся минимальный score (наихудший = наиболее аномальный).

9.3 Smoothing

def smooth_scores(scores_arr, sigma=7):
    for s in range(len(scores_arr)):
        for sig in range(1, 7):  # sigma = 1, 2, 3, 4, 5, 6
            scores_arr[s] = gaussian_filter1d(scores_arr[s], sigma=sig)
    return scores_arr
# Последовательное 6-кратное сглаживание Gaussian фильтром

10. Полный Forward Pass: Псевдокод для ONNX/torch.compile

Ниже — чистый, развёрнутый forward pass без условных ветвлений, совместимый с ONNX и torch.compile():
# ============================================================
# ONNX-Compatible STG-NF Forward Pass Pseudocode
# Input:  x: Tensor[B, 2, 24, 18] float32
# Output: nll: Tensor[B] float32
# ============================================================

def forward_onnx(x: Tensor[B, 2, 24, 18]) -> Tensor[B]:
    C, T, V = 2, 24, 18
    z = x  # (B, 2, 24, 18)
    logdet = torch.zeros(B, device=x.device)  # (B,)
    
    # ─── 8 FlowSteps ───
    for k in range(8):
        # ── 1. ActNorm ──
        z = z + actnorm_bias[k]         # (B, 2, 24, 18) + (1, 2, 1, 1)
        z = z * exp(actnorm_logs[k])     # (B, 2, 24, 18) * (1, 2, 1, 1)
        logdet = logdet + sum(actnorm_logs[k]) * T * V   # scalar * 432
        
        # ── 2. Permute (reverse) ──
        z = z[:, [1, 0], :, :]           # swap channels
        
        # ── 3. Affine Coupling ──
        z1 = z[:, 0:1, :, :]             # (B, 1, 24, 18)
        z2 = z[:, 1:2, :, :]             # (B, 1, 24, 18)
        
        # ST-GCN on z1 → h
        h = z1.clone()                    # (B, 1, 24, 18)
        
        #   GCN: Conv2d(1→2, 1×1) + einsum with A
        h = gcn_conv[k](h)               # (B, 2, 24, 18)
        h = h.view(B, 1, 2, T, V)        # (B, 1, 2, 24, 18)
        h = torch.einsum('nkctv,kvw->nctw', h, A[k])  # (B, 2, 24, 18)
        # ИЛИ эквивалент: h = h.squeeze(1); h = h @ A[k][0]
        
        #   Residual
        if k == 0:
            res = 0
        else:
            res = residual_conv[k](z1)    # Conv2d(1→2, 1×1) + BN
        
        #   TCN
        h = tcn_bn1[k](h)                # BN2d(2)
        h = relu(h)
        h = tcn_conv[k](h)               # Conv2d(2→2, (13,1), pad=(6,0))
        h = tcn_bn2[k](h)                # BN2d(2)
        h = h + res
        h = relu(h)                       # (B, 2, 24, 18)
        
        # Split h → shift, scale (cross pattern)
        shift = h[:, 0::2, :, :]          # (B, 1, 24, 18)
        scale = h[:, 1::2, :, :]          # (B, 1, 24, 18)
        
        # Squeeze dims for 3D
        shift = shift.squeeze(1)          # (B, 24, 18)
        scale = scale.squeeze(1)          # (B, 24, 18)
        z2_sq = z2.squeeze(1)             # (B, 24, 18)
        
        # Affine transform
        scale = torch.sigmoid(scale + 2.0) + 1e-6   # (B, 24, 18)
        z2_sq = (z2_sq + shift) * scale
        logdet = logdet + torch.sum(torch.log(scale), dim=[1, 2])
        
        # Reconstruct
        z1_sq = z1.squeeze(1)             # (B, 24, 18)
        z = torch.stack([z1_sq, z2_sq], dim=1)  # (B, 2, 24, 18)
    
    # ─── Prior (μ=3.0, σ²=1) ───
    mean = torch.full_like(z, 3.0)        # (B, 2, 24, 18)
    logs = torch.zeros_like(z)            # (B, 2, 24, 18)
    
    # gaussian_likelihood
    log_p = -0.5 * (logs * 2.0 + ((z - mean)**2) / (exp(logs * 2.0) + 1e-6) + log(2π))
    # При logs=0: log_p = -0.5 * ((z - 3.0)^2 + log(2π))
    likelihood = torch.sum(log_p, dim=[1, 2, 3])  # (B,)
    
    objective = logdet + likelihood
    nll = (-objective) / (math.log(2.0) * C * T * V)
    # = (-objective) / 598.449...
    
    return nll  # (B,) — чем ниже, тем более аномальный

    # Normality score = -nll (sign flip in test())

11. Таблица Тензорных Размерностей через Всю Модель


Стадия
Тензор
Shape
Dtype


Вход модели
x
(B, 2, 24, 18)
float32


logdet init
logdet
(B,)
float32


FlowStep


После ActNorm
z
(B, 2, 24, 18)
float32


После Permute
z
(B, 2, 24, 18)
float32


z1 (split)
z1
(B, 24, 18) → unsqueeze → (B, 1, 24, 18)
float32


z2 (split)
z2
(B, 24, 18) → unsqueeze → (B, 1, 24, 18)
float32


h = z1.clone()
h
(B, 1, 24, 18)
float32


gcn.conv(h)
h
(B, 2, 24, 18)
float32


gcn.view
h
(B, 1, 2, 24, 18)
float32


einsum(h, A)
h
(B, 2, 24, 18)
float32


tcn(h)
h
(B, 2, 24, 18)
float32


shift (cross)
shift
(B, 24, 18)
float32


scale (cross)
scale
(B, 24, 18)
float32


scale (sigmoid)
scale
(B, 24, 18)
float32, range (1e-6, ~1)


z2 affine
z2
(B, 24, 18)
float32


z concat
z
(B, 2, 24, 18)
float32


Prior


mean
mean
(B, 2, 24, 18)
float32, =3.0


logs
logs
(B, 2, 24, 18)
float32, =0.0


likelihood
likelihood
(B,)
float32


NLL


objective
objective
(B,)
float32


nll out
nll
(B,)
float32


12. Adjacency Matrix A — Точные Значения

При uniform, max_hops=8, openpose layout (18 joints):
# Все 18 joints связаны (max distance в скелете < 8)
# adjacency[i,j] = 1 для всех пар
# normalize_digraph: AD[i,j] = A[i,j] / sum(A[:, j])
# Каждый столбец суммируется до 1
# Т.к. все связаны: AD[i,j] = 1/18 для всех i,j (с self-loops)

A = np.ones((1, 18, 18)) / 18.0  # uniform normalized
# A[0] — матрица 18×18, все элементы = 1/18 ≈ 0.0556
Tip
Оптимизация: операция x @ A[0] где A[0] = ones(18,18)/18 это Global Average Pooling по V-измерению с broadcast обратно:
# x @ (ones/18) = mean(x, dim=-1, keepdim=True).expand(*, 18)
result = x.mean(dim=-1, keepdim=True).expand_as(x)
Это в ~18× быстрее и полностью совместимо с ONNX.


13. Рекомендации для ONNX-совместимой Реимплементации

13.1 Проблемные Паттерны для ONNX


Проблема
Расположение
Решение


torch.einsum
ConvTemporalGraphical.forward
Заменить на torch.matmul


Условный if not self.inited
_ActNorm.forward
Гарантировать inited=True перед экспортом


Lambda residual
st_gcn.__init__
Заменить на nn.Identity() или явный модуль


.clone()
FlowStep.normal_flow
Можно убрать если нет in-place ops


Dynamic squeeze/unsqueeze
split_feature, FlowStep
Зафиксировать для C=2


flow_permutation как lambda
FlowStep.__init__
Сделать обычным методом


Permute2d.indices не buffer
Permute2d
Зарегистрировать как buffer


13.2 Рекомендации для torch.compile


Аспект
Рекомендация


Data-dependent control flow
Убрать if not self.inited — поставить guard


Dynamic shapes
Зафиксировать B, C, T, V


In-place operations
Заменить inplace=True в ReLU на inplace=False


list iteration
Развернуть цикл по FlowSteps или использовать nn.Sequential


.contiguous()
Может быть лишним — проверить


13.3 Минимальный ONNX-совместимый Forward

class STG_NF_ONNX(nn.Module):
    """Упрощённая версия для экспорта в ONNX."""
    
    def __init__(self, flow_steps: nn.ModuleList):
        super().__init__()
        self.flow_steps = flow_steps  # 8 FlowStep modules
        self.R = 3.0
        self.C = 2
        self.T = 24
        self.V = 18
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args: x — (B, 2, 24, 18) float32
        Returns: nll — (B,) float32
        """
        z = x
        logdet = torch.zeros(x.shape[0], device=x.device, dtype=x.dtype)
        
        for step in self.flow_steps:
            z, logdet = step(z, logdet, reverse=False)
        
        # Prior: N(3.0, I)
        log_p = -0.5 * ((z - self.R) ** 2 + 1.8378770664093453)
        likelihood = torch.sum(log_p, dim=[1, 2, 3])
        
        objective = logdet + likelihood
        nll = (-objective) / (0.6931471805599453 * self.C * self.T * self.V)
        return nll

14. Диаграмма Потока Данных End-to-End


      flowchart TB
    subgraph Preprocessing ["Предобработка (Offline)"]
        V["Видео кадры"] --> AP["AlphaPose + YOLOX<br/>Pose Estimation"]
        AP --> PF["PoseFlow<br/>Tracking"]
        PF --> JSON["JSON: tracked_person<br/>{person_id: {frame: {kp, score}}}"]
    end
    
    subgraph DataPipeline ["Data Pipeline (Runtime)"]
        JSON --> SEG["Сегментация<br/>sliding window<br/>seg_len=24, stride=6/1"]
        SEG --> KP18["17→18 keypoints<br/>+neck"]
        KP18 --> TRANS["Transpose<br/>(N,24,18,3)→(N,3,24,18)"]
        TRANS --> SLICE["Убрать confidence<br/>[:, :2] → (N, 2, 24, 18)"]
        SLICE --> NORM["Normalize<br/>÷video_res, center, ÷std_y"]
        NORM --> AUG["Augmentation<br/>(train: +flip)"]
    end
    
    subgraph Model ["STG-NF Model"]
        AUG --> FS["8× FlowStep<br/>ActNorm → Permute → Affine(ST-GCN)"]
        FS --> Z["Latent z<br/>(B, 2, 24, 18)"]
        Z --> NLL["NLL = -(logdet + log p(z))<br/>/ (log2 · 2 · 24 · 18)"]
    end
    
    subgraph Scoring ["Scoring"]
        NLL --> FRS["Frame Score<br/>= -NLL per segment<br/>assigned to frame t+12"]
        FRS --> MIN["Min over persons<br/>per frame"]
        MIN --> SMOOTH["6× Gaussian smooth<br/>σ = 1..6"]
        SMOOTH --> AUC["ROC-AUC vs GT"]
    end

    
      Loading

  
15. Скелетный Граф — Визуализация

OpenPose 18 Keypoints (COCO18 reordered):
 0: Nose
 1: Neck  (synthetic = avg shoulders)
 2: RShoulder
 3: RElbow
 4: RWrist
 5: LShoulder
 6: LElbow
 7: LWrist
 8: RHip
 9: RKnee
10: RAnkle
11: LHip
12: LKnee
13: LAnkle
14: REye
15: LEye
16: REar
17: LEar

Рёбра (anatomical):
  0─1 (nose─neck)
  1─2, 1─5 (neck─shoulders)
  2─3─4 (R arm)
  5─6─7 (L arm)
  2─8, 5─11 (torso)
  8─9─10 (R leg)
  11─12─13 (L leg)
  0─14, 0─15 (eyes)
  14─16, 15─17 (ears)


16. Важные Нюансы и Edge Cases

16.1 Split Feature Squeeze Behavior

def split_feature(tensor, type="split"):
    C = tensor.size(1)  # = 2
    if type == "split":
        # tensor[:, :1, ...] → (B, 1, 24, 18).squeeze(dim=1) → (B, 24, 18) [3D!]
        return tensor[:, :C//2, ...].squeeze(dim=1), tensor[:, C//2:, ...].squeeze(dim=1)
    elif type == "cross":
        # tensor[:, 0::2, ...] → (B, 1, 24, 18).squeeze(dim=1) → (B, 24, 18)
        return tensor[:, 0::2, ...].squeeze(dim=1), tensor[:, 1::2, ...].squeeze(dim=1)
Критично: После split, z1 и z2 — 3D (B, 24, 18). Код проверяет if len(z1.shape) == 3 и делает unsqueeze(dim=1). Для ONNX нужно это зафиксировать.
16.2 Edge Importance

При edge_importance=False: self.edge_importance = [1] * len(self.block) — Python list, не tensor. В цикле A * importance = A * 1 = A. Для ONNX это прозрачно.
16.3 Gradient Clipping

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=100)
16.4 Score Sign Convention


Модель выдаёт NLL: чем ВЫШЕ nll, тем более аномальный
Test: probs = -1 * nll — инверсия: чем ВЫШЕ prob, тем более нормальный
ROC-AUC: GT: 1=normal, 0=abnormal. Scores: higher = more normal → корректный AUC

16.5 Smoothing Details

# GT маски ShanghaiTech: загружается .npy и инвертируется
clip_gt = np.ones(shape) - clip_gt  # 0→1(normal), 1→0(abnormal)

# Score assignment: pid_frame_inds + seg_len/2 = frame + 12
# Frames without predictions get score = inf
# inf заменяется на max(non-inf scores)

17. Полная Спецификация для Воссоздания Модели

Checklist для имплементации:


Graph: Создать полносвязную матрицу 18×18, нормализовать по столбцам → (1, 18, 18) float32
8 FlowSteps: каждый = ActNorm(2) + Permute([1,0]) + AffineCoupling(ST-GCN)
ST-GCN в coupling: Conv2d(1→2, 1×1) @ A + TCN(BN+ReLU+Conv(13×1)+BN) + residual + ReLU
Residual: k=0: None; k=1-7: Conv2d(1→2, 1×1)+BN2d
Affine: shift, scale = cross_split(stgcn_output); scale = sigmoid(scale+2)+1e-6; z2 = (z2+shift)*scale
Prior: μ=3.0, σ²=1 (logs=0)
NLL: -(logdet + gaussian_log_likelihood) / (log2 * 2 * 24 * 18)
Scoring: -NLL per segment, assign to mid-frame, min over persons, 6× Gaussian smooth, ROC-AUC
Параметр	Значение	Описание
dataset	`ShanghaiTech`	Датасет
`seg_len` (τ)	`24`	Длина сегмента позы (кадры)
`seg_stride`	`6`	Шаг скользящего окна (train). Test = `1`
`K`	`8`	Количество Flow Steps
`L`	`1`	Количество уровней (без Squeeze)
`R`	`3.0`	Среднее prior-распределения μ_normal
`hidden_channels`	`0`	Скрытые каналы ST-GCN (0 = один блок)
`flow_permutation`	`permute`	Тип перестановки (reverse order)
`flow_coupling`	`affine`	Тип coupling layer
`adj_strategy`	`uniform`	Стратегия графа
`max_hops`	`8`	Макс. дистанция в adjacency
`edge_importance`	`False`	Веса рёбер
`temporal_kernel`	`None` → `13`	Авто: `T//2 + 1 = 24//2+1 = 13`
`model_confidence`	`False`	Не использовать канал confidence
`headless`	`False`	Все 18 keypoints
`batch_size`	`256`	Размер батча
`epochs`	`8`	Число эпох
`model_lr`	`5e-4`	Learning rate
`model_optimizer`	`adamx`	Optimax optimizer
`model_sched`	`exp_decay` → `None`	Scheduler = None (ручной decay)
`model_lr_decay`	`0.99`	LR decay per epoch
`model_weight_decay`	`5e-5`	Weight decay
`LU_decomposed`	`True`	LU-разложение для InvConv
`learn_top`	`False`	Без обучаемого prior
`actnorm_scale`	`1.0`	Масштаб ActNorm
`num_transform`	`2`	Число аугментаций (identity + flip)
`prop_norm_scale`	`1`	Пропорциональная нормализация
`norm_scale`	`0`	Не масштабировать непропорционально
`global_pose_segs`	`True`	Нормализовать позы
`seed`	`999`	Random seed
Стадия	Shape	Dtype	Описание
Raw JSON keypoints	(51,) flat	float	`[x1,y1,c1,...,x17,y17,c17]`
Single person track	(T_clip, 17, 3)	float64	Все кадры 1 человека
After segmentation	(N_seg, 24, 17, 3)	float64	Скользящее окно
After 17→18	(N_seg, 24, 18, 3)	float64	+neck keypoint
After transpose	(N_seg, 3, 24, 18) = (N, C, T, V)	float32	Каналы первые
After `[:, :2]`	(B, 2, 24, 18)	float32	Без confidence
After normalize	(B, 2, 24, 18)	float32	Вход модели
FlowStep k	residual
k=0 (first)	`lambda x: 0` — нет residual
k=1..7	Conv2d(1→2, 1×1) + BN2d(2) — проекция каналов
Компонент	Параметры	Shape	Count
ActNorm
bias	learnable	(1, 2, 1, 1)	2
logs	learnable	(1, 2, 1, 1)	2
Permute2d	—	indices only	0
ST-GCN
gcn.conv.weight	learnable	(2, 1, 1, 1)	2
gcn.conv.bias	learnable	(2,)	2
tcn.BN1.weight	learnable	(2,)	2
tcn.BN1.bias	learnable	(2,)	2
tcn.Conv2d.weight	learnable	(2, 2, 13, 1)	52
tcn.Conv2d.bias	learnable	(2,)	2
tcn.BN2.weight	learnable	(2,)	2
tcn.BN2.bias	learnable	(2,)	2
Residual	None (=0)		0
Subtotal			68
Компонент	Параметры	Shape	Count
ActNorm	bias + logs	2 × (1,2,1,1)	4
Permute2d	—		0
ST-GCN (same as k=0)			64
Residual Conv2d	weight	(2, 1, 1, 1)	2
Residual Conv2d	bias	—	(no bias, BN follows)
Residual BN2d	weight+bias	(2,) × 2	4
Subtotal			74
Стадия	Тензор	Shape	Dtype
Вход модели	`x`	(B, 2, 24, 18)	float32
logdet init	`logdet`	(B,)	float32
FlowStep
После ActNorm	`z`	(B, 2, 24, 18)	float32
После Permute	`z`	(B, 2, 24, 18)	float32
z1 (split)	`z1`	(B, 24, 18) → unsqueeze → (B, 1, 24, 18)	float32
z2 (split)	`z2`	(B, 24, 18) → unsqueeze → (B, 1, 24, 18)	float32
h = z1.clone()	`h`	(B, 1, 24, 18)	float32
gcn.conv(h)	`h`	(B, 2, 24, 18)	float32
gcn.view	`h`	(B, 1, 2, 24, 18)	float32
einsum(h, A)	`h`	(B, 2, 24, 18)	float32
tcn(h)	`h`	(B, 2, 24, 18)	float32
shift (cross)	`shift`	(B, 24, 18)	float32
scale (cross)	scale	(B, 24, 18)	float32
scale (sigmoid)	scale	(B, 24, 18)	float32, range (1e-6, ~1)
z2 affine	`z2`	(B, 24, 18)	float32
z concat	`z`	(B, 2, 24, 18)	float32
Prior
mean	`mean`	(B, 2, 24, 18)	float32, =3.0
logs	`logs`	(B, 2, 24, 18)	float32, =0.0
likelihood	likelihood	(B,)	float32
NLL
objective	`objective`	(B,)	float32
nll out	`nll`	(B,)	float32
Проблема	Расположение	Решение
`torch.einsum`	`ConvTemporalGraphical.forward`	Заменить на `torch.matmul`
Условный `if not self.inited`	`_ActNorm.forward`	Гарантировать `inited=True` перед экспортом
Lambda residual	`st_gcn.__init__`	Заменить на `nn.Identity()` или явный модуль
`.clone()`	`FlowStep.normal_flow`	Можно убрать если нет in-place ops
Dynamic squeeze/unsqueeze	split_feature, FlowStep	Зафиксировать для C=2
`flow_permutation` как lambda	`FlowStep.__init__`	Сделать обычным методом
`Permute2d.indices` не buffer	Permute2d	Зарегистрировать как buffer
Аспект	Рекомендация
Data-dependent control flow	Убрать `if not self.inited` — поставить guard
Dynamic shapes	Зафиксировать B, C, T, V
In-place operations	Заменить `inplace=True` в ReLU на `inplace=False`
list iteration	Развернуть цикл по FlowSteps или использовать `nn.Sequential`
`.contiguous()`	Может быть лишним — проверить