mirror of
https://github.com/ruvnet/RuView.git
synced 2026-04-28 05:59:32 +00:00
327d0d13f6
Add --scale flag with 4 presets for dataset-appropriate sizing: lite: ~190K params, 2 TCN blocks k=3 (trains in seconds) small: ~200K params, 4 TCN blocks k=5 (trains in minutes) medium: ~800K params, 4 TCN blocks k=7 (trains in ~15 min) full: ~7.7M params, 4 TCN blocks k=7 (trains in hours) Refactored model to use dynamic TCN block count, kernel size, channel widths, hidden dim, and SPSA perturbation count — all driven by the scale preset. Default is 'lite' for fast iteration. Validated: lite model completes 30 epochs on 265 samples in ~2 min on Windows CPU (vs stuck at epoch 1 with full model). Scale up with: --scale small|medium|full as dataset grows. Co-Authored-By: claude-flow <ruv@ruv.net>
1657 lines
55 KiB
JavaScript
1657 lines
55 KiB
JavaScript
#!/usr/bin/env node
|
|
/**
|
|
* WiFlow Supervised Pose Training Pipeline (ADR-079)
|
|
*
|
|
* Trains WiFlow pose estimation on paired CSI + camera keypoint data.
|
|
* Extends the ruvllm training infrastructure with a simplified TCN architecture
|
|
* and three-phase curriculum: contrastive pretraining, supervised keypoint
|
|
* regression, and refinement with bone/temporal constraints.
|
|
*
|
|
* Input format (paired JSONL):
|
|
* {"csi": [[...128 or 8 floats...], ...20 frames], "keypoints": [[x,y],...17], "conf": [c0..c16], "timestamp": ...}
|
|
*
|
|
* Architecture:
|
|
* TCN (4 dilated causal conv blocks, k=7, dilation 1,2,4,8)
|
|
* input_dim -> 256 -> 192 -> 128
|
|
* Flatten [128*20] -> Linear 2560 -> 2048 -> Linear 2048 -> 34
|
|
* Reshape to [17, 2] keypoints in [0, 1]
|
|
*
|
|
* Phases:
|
|
* 1. Contrastive (50 epochs) — representation learning on CSI windows
|
|
* 2. Supervised (200 epochs) — confidence-weighted SmoothL1 on keypoints
|
|
* with curriculum: conf>0.9 -> conf>0.7 -> conf>0.5 -> all + augmentation
|
|
* 3. Refinement (50 epochs) — combined loss with bone + temporal constraints
|
|
*
|
|
* Usage:
|
|
* node scripts/train-wiflow-supervised.js --data data/paired-csi-keypoints.jsonl
|
|
* node scripts/train-wiflow-supervised.js --data data/paired.jsonl --skip-contrastive --epochs 200
|
|
* node scripts/train-wiflow-supervised.js --data data/paired.jsonl --output models/wiflow-sup-v2
|
|
*
|
|
* ADR: docs/adr/ADR-079
|
|
*/
|
|
|
|
'use strict';
|
|
|
|
const fs = require('fs');
|
|
const path = require('path');
|
|
const { parseArgs } = require('util');
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Resolve ruvllm from vendor tree
|
|
// ---------------------------------------------------------------------------
|
|
const RUVLLM_PATH = path.resolve(__dirname, '..', 'vendor', 'ruvector', 'npm', 'packages', 'ruvllm', 'src');
|
|
|
|
const {
|
|
ContrastiveTrainer,
|
|
cosineSimilarity,
|
|
infoNCELoss,
|
|
computeGradient,
|
|
} = require(path.join(RUVLLM_PATH, 'contrastive.js'));
|
|
|
|
const {
|
|
TrainingPipeline,
|
|
} = require(path.join(RUVLLM_PATH, 'training.js'));
|
|
|
|
const {
|
|
EwcManager,
|
|
} = require(path.join(RUVLLM_PATH, 'sona.js'));
|
|
|
|
const {
|
|
SafeTensorsWriter,
|
|
ModelExporter,
|
|
} = require(path.join(RUVLLM_PATH, 'export.js'));
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CLI argument parsing
|
|
// ---------------------------------------------------------------------------
|
|
const { values: args } = parseArgs({
|
|
options: {
|
|
data: { type: 'string', short: 'd' },
|
|
output: { type: 'string', short: 'o', default: 'models/wiflow-supervised' },
|
|
epochs: { type: 'string', short: 'e', default: '300' },
|
|
'batch-size': { type: 'string', default: '32' },
|
|
lr: { type: 'string', default: '0.0001' },
|
|
'skip-contrastive': { type: 'boolean', default: false },
|
|
'eval-split': { type: 'string', default: '0.2' },
|
|
scale: { type: 'string', short: 's', default: 'lite' },
|
|
verbose: { type: 'boolean', short: 'v', default: false },
|
|
},
|
|
strict: true,
|
|
});
|
|
|
|
if (!args.data) {
|
|
console.error('Usage: node scripts/train-wiflow-supervised.js --data <paired-jsonl> [options]');
|
|
console.error('');
|
|
console.error('Options:');
|
|
console.error(' --data <file> Paired CSI+keypoint JSONL (required)');
|
|
console.error(' --output <dir> Output directory (default: models/wiflow-supervised)');
|
|
console.error(' --epochs <n> Total epochs across all phases (default: 300)');
|
|
console.error(' --batch-size <n> Batch size (default: 32)');
|
|
console.error(' --lr <float> Learning rate (default: 0.0001)');
|
|
console.error(' --skip-contrastive Skip phase 1 contrastive pretraining');
|
|
console.error(' --eval-split <float> Held-out eval fraction (default: 0.2)');
|
|
console.error(' --verbose Print detailed progress');
|
|
process.exit(1);
|
|
}
|
|
|
|
const CONFIG = {
|
|
dataPath: args.data,
|
|
outputDir: args.output,
|
|
totalEpochs: parseInt(args.epochs, 10),
|
|
batchSize: parseInt(args['batch-size'], 10),
|
|
lr: parseFloat(args.lr),
|
|
skipContrastive: args['skip-contrastive'],
|
|
evalSplit: parseFloat(args['eval-split']),
|
|
verbose: args.verbose,
|
|
|
|
// Phase epoch allocation (scaled to totalEpochs)
|
|
contrastiveRatio: 50 / 300,
|
|
supervisedRatio: 200 / 300,
|
|
refinementRatio: 50 / 300,
|
|
|
|
// Curriculum confidence thresholds (O1)
|
|
curriculumStages: [0.9, 0.7, 0.5, 0.0],
|
|
|
|
// Architecture
|
|
timeSteps: 20,
|
|
numKeypoints: 17,
|
|
|
|
// SGD momentum
|
|
momentum: 0.9,
|
|
|
|
// Refinement loss weights
|
|
boneWeight: 0.3,
|
|
temporalWeight: 0.1,
|
|
};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Model scale presets: lite → small → medium → full
|
|
// lite: ~45K params, trains in seconds (good for <1K samples)
|
|
// small: ~200K params, trains in minutes (good for 1K-10K samples)
|
|
// medium: ~800K params, trains in ~15 min (good for 10K-50K samples)
|
|
// full: ~7.7M params, trains in hours (good for 50K+ samples)
|
|
// ---------------------------------------------------------------------------
|
|
const SCALE_PRESETS = {
|
|
lite: { tcnChannels: [32, 32, 32, 32], hiddenDim: 256, tcnBlocks: 2, kernel: 3, spsaK: 1 },
|
|
small: { tcnChannels: [64, 64, 48, 32], hiddenDim: 512, tcnBlocks: 4, kernel: 5, spsaK: 2 },
|
|
medium: { tcnChannels: [128, 128, 96, 64], hiddenDim: 1024, tcnBlocks: 4, kernel: 7, spsaK: 3 },
|
|
full: { tcnChannels: [256, 256, 192, 128], hiddenDim: 2048, tcnBlocks: 4, kernel: 7, spsaK: 3 },
|
|
};
|
|
|
|
const scaleKey = args.scale || 'lite';
|
|
const SCALE = SCALE_PRESETS[scaleKey] || SCALE_PRESETS.lite;
|
|
console.log(`Model scale: ${scaleKey} (${JSON.stringify(SCALE)})`);
|
|
|
|
// Compute phase epochs
|
|
const totalForPhases = CONFIG.skipContrastive
|
|
? CONFIG.totalEpochs
|
|
: CONFIG.totalEpochs;
|
|
const contrastiveEpochs = CONFIG.skipContrastive ? 0 : Math.round(totalForPhases * CONFIG.contrastiveRatio);
|
|
const supervisedEpochs = Math.round(totalForPhases * CONFIG.supervisedRatio);
|
|
const refinementEpochs = totalForPhases - contrastiveEpochs - supervisedEpochs;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Deterministic PRNG (xorshift32)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function createRng(seed) {
|
|
let s = seed | 0 || 42;
|
|
return () => {
|
|
s ^= s << 13;
|
|
s ^= s >> 17;
|
|
s ^= s << 5;
|
|
return (s >>> 0) / 4294967296;
|
|
};
|
|
}
|
|
|
|
function gaussianRng(rng) {
|
|
return () => {
|
|
const u1 = rng() || 1e-10;
|
|
const u2 = rng();
|
|
return Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// O6: Subcarrier importance scoring (ruvector-solver inspired)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Score each subcarrier by temporal variance — high-variance subcarriers
|
|
* carry motion information, low-variance ones are noise/static.
|
|
* Returns sorted indices of top-K most informative subcarriers.
|
|
* This is the JS equivalent of ruvector-solver's sparse interpolation (114→56).
|
|
*/
|
|
function selectTopSubcarriers(samples, dim, T, topK) {
|
|
const variance = new Float64Array(dim);
|
|
for (const s of samples) {
|
|
for (let d = 0; d < dim; d++) {
|
|
let mean = 0;
|
|
for (let t = 0; t < T; t++) mean += s.csi[d * T + t];
|
|
mean /= T;
|
|
let v = 0;
|
|
for (let t = 0; t < T; t++) v += (s.csi[d * T + t] - mean) ** 2;
|
|
variance[d] += v / T;
|
|
}
|
|
}
|
|
// Average variance across samples
|
|
for (let d = 0; d < dim; d++) variance[d] /= samples.length;
|
|
|
|
// Rank by variance (descending)
|
|
const indices = Array.from({ length: dim }, (_, i) => i);
|
|
indices.sort((a, b) => variance[b] - variance[a]);
|
|
return indices.slice(0, topK);
|
|
}
|
|
|
|
/**
|
|
* Reduce CSI samples to selected subcarrier indices.
|
|
* [dim, T] → [topK, T]
|
|
*/
|
|
function reduceSubcarriers(sample, selectedIndices, T) {
|
|
const topK = selectedIndices.length;
|
|
const reduced = new Float32Array(topK * T);
|
|
for (let k = 0; k < topK; k++) {
|
|
const srcD = selectedIndices[k];
|
|
for (let t = 0; t < T; t++) {
|
|
reduced[k * T + t] = sample.csi[srcD * T + t];
|
|
}
|
|
}
|
|
return { ...sample, csi: reduced, csiDim: topK };
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// O7: Attention-weighted subcarrier scoring (ruvector-attention inspired)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Compute spatial attention weights for subcarriers based on correlation
|
|
* with ground-truth keypoint motion. Subcarriers that covary with skeleton
|
|
* movement get higher weight.
|
|
* Returns Float32Array[dim] of attention weights (sum = 1).
|
|
*/
|
|
function computeSubcarrierAttention(samples, dim, T) {
|
|
const weights = new Float64Array(dim);
|
|
|
|
for (const s of samples) {
|
|
// Compute per-subcarrier energy (proxy for motion sensitivity)
|
|
for (let d = 0; d < dim; d++) {
|
|
let energy = 0;
|
|
for (let t = 1; t < T; t++) {
|
|
const diff = s.csi[d * T + t] - s.csi[d * T + (t - 1)];
|
|
energy += diff * diff;
|
|
}
|
|
// Weight by confidence — higher confidence samples matter more
|
|
const confWeight = s.conf ? (s.conf.reduce((a, b) => a + b, 0) / s.conf.length) : 1.0;
|
|
weights[d] += energy * confWeight;
|
|
}
|
|
}
|
|
|
|
// Softmax normalization
|
|
let maxW = -Infinity;
|
|
for (let d = 0; d < dim; d++) if (weights[d] > maxW) maxW = weights[d];
|
|
let sumExp = 0;
|
|
const attn = new Float32Array(dim);
|
|
for (let d = 0; d < dim; d++) {
|
|
attn[d] = Math.exp((weights[d] - maxW) / (maxW * 0.1 + 1e-8)); // temperature scaling
|
|
sumExp += attn[d];
|
|
}
|
|
for (let d = 0; d < dim; d++) attn[d] /= sumExp;
|
|
|
|
return attn;
|
|
}
|
|
|
|
/**
|
|
* Apply attention weights to CSI input: weight each subcarrier channel.
|
|
*/
|
|
function applySubcarrierAttention(csi, attn, dim, T) {
|
|
const weighted = new Float32Array(csi.length);
|
|
for (let d = 0; d < dim; d++) {
|
|
const w = attn[d] * dim; // Rescale so mean weight = 1
|
|
for (let t = 0; t < T; t++) {
|
|
weighted[d * T + t] = csi[d * T + t] * w;
|
|
}
|
|
}
|
|
return weighted;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// O8: DynamicMinCut multi-person separation (ruvector-mincut inspired)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* JS implementation of Stoer-Wagner min-cut for person separation in CSI.
|
|
* Builds a correlation graph where subcarriers are nodes and edges are
|
|
* temporal correlation. Min-cut separates subcarrier groups that respond
|
|
* to different people.
|
|
*
|
|
* Returns partition assignments [0 or 1] per subcarrier.
|
|
*/
|
|
function stoerWagnerMinCut(adjacency, n) {
|
|
// Stoer-Wagner: find global min-cut by repeated minimum-cut-phase
|
|
let bestCut = Infinity;
|
|
let bestPartition = null;
|
|
|
|
// Work on a copy with merged-node tracking
|
|
const merged = new Array(n).fill(false);
|
|
const adj = [];
|
|
for (let i = 0; i < n; i++) {
|
|
adj[i] = new Float64Array(n);
|
|
for (let j = 0; j < n; j++) adj[i][j] = adjacency[i * n + j];
|
|
}
|
|
const nodeMap = Array.from({ length: n }, (_, i) => [i]); // track merged nodes
|
|
|
|
for (let phase = 0; phase < n - 1; phase++) {
|
|
// Minimum cut phase
|
|
const inA = new Array(n).fill(false);
|
|
const w = new Float64Array(n); // connectivity to set A
|
|
let last = -1, secondLast = -1;
|
|
|
|
for (let step = 0; step < n - phase; step++) {
|
|
// Find most tightly connected vertex not in A
|
|
let maxW = -1, maxIdx = -1;
|
|
for (let v = 0; v < n; v++) {
|
|
if (!merged[v] && !inA[v] && w[v] > maxW) {
|
|
maxW = w[v];
|
|
maxIdx = v;
|
|
}
|
|
}
|
|
if (maxIdx === -1) {
|
|
// Find any unmerged non-A vertex
|
|
for (let v = 0; v < n; v++) {
|
|
if (!merged[v] && !inA[v]) { maxIdx = v; break; }
|
|
}
|
|
}
|
|
if (maxIdx === -1) break;
|
|
|
|
secondLast = last;
|
|
last = maxIdx;
|
|
inA[maxIdx] = true;
|
|
|
|
// Update weights
|
|
for (let v = 0; v < n; v++) {
|
|
if (!merged[v] && !inA[v]) {
|
|
w[v] += adj[maxIdx][v];
|
|
}
|
|
}
|
|
}
|
|
|
|
if (last === -1 || secondLast === -1) break;
|
|
|
|
// Cut of the phase = w[last]
|
|
const cutVal = w[last];
|
|
if (cutVal < bestCut) {
|
|
bestCut = cutVal;
|
|
bestPartition = new Array(n).fill(0);
|
|
for (const idx of nodeMap[last]) bestPartition[idx] = 1;
|
|
}
|
|
|
|
// Merge last into secondLast
|
|
for (let v = 0; v < n; v++) {
|
|
adj[secondLast][v] += adj[last][v];
|
|
adj[v][secondLast] += adj[v][last];
|
|
}
|
|
adj[secondLast][secondLast] = 0;
|
|
nodeMap[secondLast] = nodeMap[secondLast].concat(nodeMap[last]);
|
|
merged[last] = true;
|
|
}
|
|
|
|
return { cutValue: bestCut, partition: bestPartition || new Array(n).fill(0) };
|
|
}
|
|
|
|
/**
|
|
* Build subcarrier correlation graph and apply min-cut to separate
|
|
* person-specific subcarrier clusters.
|
|
* Returns: { partition: [0|1 per subcarrier], cutValue: float }
|
|
*/
|
|
function minCutPersonSeparation(samples, dim, T) {
|
|
// Build correlation matrix across subcarriers
|
|
const corr = new Float64Array(dim * dim);
|
|
|
|
for (const s of samples) {
|
|
for (let i = 0; i < dim; i++) {
|
|
for (let j = i + 1; j < dim; j++) {
|
|
// Pearson correlation between subcarrier i and j
|
|
let sumI = 0, sumJ = 0, sumIJ = 0, sumI2 = 0, sumJ2 = 0;
|
|
for (let t = 0; t < T; t++) {
|
|
const vi = s.csi[i * T + t];
|
|
const vj = s.csi[j * T + t];
|
|
sumI += vi; sumJ += vj;
|
|
sumIJ += vi * vj;
|
|
sumI2 += vi * vi; sumJ2 += vj * vj;
|
|
}
|
|
const num = T * sumIJ - sumI * sumJ;
|
|
const den = Math.sqrt((T * sumI2 - sumI * sumI) * (T * sumJ2 - sumJ * sumJ));
|
|
const r = den > 1e-8 ? Math.abs(num / den) : 0;
|
|
corr[i * dim + j] = r;
|
|
corr[j * dim + i] = r;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Average across samples
|
|
const nSamples = samples.length || 1;
|
|
for (let i = 0; i < corr.length; i++) corr[i] /= nSamples;
|
|
|
|
return stoerWagnerMinCut(corr, dim);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// O9: Multi-SPSA gradient estimation (improved convergence)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Multi-perturbation SPSA: average over K random directions per step.
|
|
* Reduces variance by sqrt(K) compared to single SPSA.
|
|
* K=3 gives 1.7x better gradient estimates at 3x forward passes (net win
|
|
* because gradient quality matters more than speed for convergence).
|
|
*/
|
|
function multiSpsaGrad(model, batch, lossFn, paramObj, rng, K) {
|
|
K = K || 3;
|
|
const eps = 1e-4;
|
|
const w = paramObj.weight;
|
|
const n = w.length;
|
|
const grad = new Float32Array(n);
|
|
|
|
for (let k = 0; k < K; k++) {
|
|
const delta = new Float32Array(n);
|
|
for (let i = 0; i < n; i++) delta[i] = rng() < 0.5 ? 1 : -1;
|
|
|
|
// w + eps*delta
|
|
for (let i = 0; i < n; i++) w[i] += eps * delta[i];
|
|
let lp = 0;
|
|
for (const s of batch) lp += lossFn(model, s);
|
|
lp /= batch.length;
|
|
|
|
// w - eps*delta
|
|
for (let i = 0; i < n; i++) w[i] -= 2 * eps * delta[i];
|
|
let lm = 0;
|
|
for (const s of batch) lm += lossFn(model, s);
|
|
lm /= batch.length;
|
|
|
|
// Restore
|
|
for (let i = 0; i < n; i++) w[i] += eps * delta[i];
|
|
|
|
const scale = (lp - lm) / (2 * eps);
|
|
for (let i = 0; i < n; i++) grad[i] += scale / delta[i];
|
|
}
|
|
|
|
// Average over K perturbations
|
|
for (let i = 0; i < n; i++) grad[i] /= K;
|
|
return grad;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tensor utilities
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function initKaiming(fanIn, fanOut, rng) {
|
|
const std = Math.sqrt(2.0 / fanIn);
|
|
const gauss = gaussianRng(rng);
|
|
const arr = new Float32Array(fanIn * fanOut);
|
|
for (let i = 0; i < arr.length; i++) arr[i] = gauss() * std;
|
|
return arr;
|
|
}
|
|
|
|
function initXavier(fanIn, fanOut, rng) {
|
|
const std = Math.sqrt(2.0 / (fanIn + fanOut));
|
|
const gauss = gaussianRng(rng);
|
|
const arr = new Float32Array(fanIn * fanOut);
|
|
for (let i = 0; i < arr.length; i++) arr[i] = gauss() * std;
|
|
return arr;
|
|
}
|
|
|
|
function relu(arr) {
|
|
for (let i = 0; i < arr.length; i++) {
|
|
if (arr[i] < 0) arr[i] = 0;
|
|
}
|
|
return arr;
|
|
}
|
|
|
|
function sigmoid(x) {
|
|
return 1.0 / (1.0 + Math.exp(-x));
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// SmoothL1 loss and gradient
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function smoothL1(predicted, target, beta) {
|
|
beta = beta || 0.05;
|
|
let loss = 0;
|
|
const n = Math.min(predicted.length, target.length);
|
|
for (let i = 0; i < n; i++) {
|
|
const diff = Math.abs(predicted[i] - target[i]);
|
|
if (diff < beta) {
|
|
loss += 0.5 * diff * diff / beta;
|
|
} else {
|
|
loss += diff - 0.5 * beta;
|
|
}
|
|
}
|
|
return loss / n;
|
|
}
|
|
|
|
function smoothL1Grad(predicted, target, beta) {
|
|
beta = beta || 0.05;
|
|
const n = Math.min(predicted.length, target.length);
|
|
const grad = new Float32Array(n);
|
|
for (let i = 0; i < n; i++) {
|
|
const diff = predicted[i] - target[i];
|
|
const absDiff = Math.abs(diff);
|
|
if (absDiff < beta) {
|
|
grad[i] = diff / beta / n;
|
|
} else {
|
|
grad[i] = (diff > 0 ? 1 : -1) / n;
|
|
}
|
|
}
|
|
return grad;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// COCO bone priors (ADR-079)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const BONE_CONNECTIONS = [
|
|
[0, 1], [0, 2], // nose -> eyes
|
|
[1, 3], [2, 4], // eyes -> ears
|
|
[5, 7], [7, 9], // left arm: shoulder-elbow, elbow-wrist
|
|
[6, 8], [8, 10], // right arm: shoulder-elbow, elbow-wrist
|
|
[5, 11], [6, 12], // torso: shoulder-hip
|
|
[11, 13], [13, 15], // left leg: hip-knee, knee-ankle
|
|
[12, 14], [14, 16], // right leg: hip-knee, knee-ankle
|
|
[5, 6], // shoulder width
|
|
];
|
|
|
|
const BONE_LENGTH_PRIORS = [
|
|
0.06, 0.06, // nose-eye
|
|
0.06, 0.06, // eye-ear
|
|
0.15, 0.13, // left shoulder-elbow, elbow-wrist
|
|
0.15, 0.13, // right shoulder-elbow, elbow-wrist
|
|
0.26, 0.26, // shoulder-hip
|
|
0.25, 0.25, // left hip-knee, knee-ankle
|
|
0.25, 0.25, // right hip-knee, knee-ankle
|
|
0.20, // shoulder width
|
|
];
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Data loading — paired CSI + keypoint JSONL
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Load paired dataset from JSONL file.
|
|
* Each line: { csi: [[...], ...], keypoints: [[x,y], ...17], conf: [...17], timestamp: ... }
|
|
* csi shape: [subcarriers, timeSteps] or [features, timeSteps]
|
|
*/
|
|
function loadPairedData(filePath) {
|
|
if (!fs.existsSync(filePath)) {
|
|
console.error(`Data file not found: ${filePath}`);
|
|
process.exit(1);
|
|
}
|
|
|
|
const content = fs.readFileSync(filePath, 'utf-8');
|
|
const lines = content.split('\n').filter(l => l.trim());
|
|
const samples = [];
|
|
|
|
for (const line of lines) {
|
|
try {
|
|
const obj = JSON.parse(line);
|
|
if (!obj.csi || !(obj.keypoints || obj.kp)) continue;
|
|
|
|
const csi = obj.csi; // 2D array [dim, T] or flat
|
|
const kp = obj.keypoints || obj.kp; // [[x,y], ...] or flat [x,y,x,y,...]
|
|
const conf = obj.conf || null; // [c0, c1, ...c16] or scalar or null
|
|
const ts = obj.timestamp || obj.ts_start || 0;
|
|
|
|
// Flatten keypoints to [34] = [x0, y0, x1, y1, ...]
|
|
let kpFlat;
|
|
if (Array.isArray(kp[0])) {
|
|
kpFlat = new Float32Array(CONFIG.numKeypoints * 2);
|
|
for (let i = 0; i < CONFIG.numKeypoints && i < kp.length; i++) {
|
|
kpFlat[i * 2] = kp[i][0];
|
|
kpFlat[i * 2 + 1] = kp[i][1];
|
|
}
|
|
} else {
|
|
kpFlat = new Float32Array(kp.slice(0, CONFIG.numKeypoints * 2));
|
|
}
|
|
|
|
// Confidence per keypoint
|
|
let confArr;
|
|
if (conf && Array.isArray(conf) && conf.length >= CONFIG.numKeypoints) {
|
|
confArr = new Float32Array(conf.slice(0, CONFIG.numKeypoints));
|
|
} else if (typeof conf === 'number') {
|
|
confArr = new Float32Array(CONFIG.numKeypoints).fill(conf);
|
|
} else {
|
|
confArr = new Float32Array(CONFIG.numKeypoints).fill(1.0);
|
|
}
|
|
|
|
// Flatten CSI to Float32Array [dim * T]
|
|
let csiFlat;
|
|
let csiDim;
|
|
if (Array.isArray(csi[0])) {
|
|
csiDim = csi.length;
|
|
const T = csi[0].length;
|
|
csiFlat = new Float32Array(csiDim * T);
|
|
for (let d = 0; d < csiDim; d++) {
|
|
for (let t = 0; t < T; t++) {
|
|
csiFlat[d * T + t] = csi[d][t] || 0;
|
|
}
|
|
}
|
|
} else if (obj.csi_shape && obj.csi_shape.length === 2) {
|
|
// Flat array with explicit shape: [dim, T]
|
|
csiDim = obj.csi_shape[0];
|
|
csiFlat = new Float32Array(csi);
|
|
} else {
|
|
csiDim = csi.length;
|
|
csiFlat = new Float32Array(csi);
|
|
}
|
|
|
|
samples.push({ csi: csiFlat, csiDim, keypoints: kpFlat, conf: confArr, timestamp: ts });
|
|
} catch (_) {
|
|
// Skip malformed lines
|
|
}
|
|
}
|
|
|
|
return samples;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Data augmentation (O2)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function augmentSample(sample, rng, T) {
|
|
const dim = sample.csiDim;
|
|
const augCsi = new Float32Array(sample.csi);
|
|
|
|
// Time shift: roll ±2 frames
|
|
const shift = Math.floor(rng() * 5) - 2; // -2 to +2
|
|
if (shift !== 0) {
|
|
const temp = new Float32Array(dim * T);
|
|
for (let d = 0; d < dim; d++) {
|
|
for (let t = 0; t < T; t++) {
|
|
let srcT = t - shift;
|
|
if (srcT < 0) srcT = 0;
|
|
if (srcT >= T) srcT = T - 1;
|
|
temp[d * T + t] = augCsi[d * T + srcT];
|
|
}
|
|
}
|
|
augCsi.set(temp);
|
|
}
|
|
|
|
// Amplitude noise: gaussian sigma=0.02
|
|
const gauss = gaussianRng(rng);
|
|
for (let i = 0; i < augCsi.length; i++) {
|
|
augCsi[i] += gauss() * 0.02;
|
|
}
|
|
|
|
// Subcarrier dropout: zero 10% randomly
|
|
for (let d = 0; d < dim; d++) {
|
|
if (rng() < 0.10) {
|
|
for (let t = 0; t < T; t++) {
|
|
augCsi[d * T + t] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
return {
|
|
csi: augCsi,
|
|
csiDim: dim,
|
|
keypoints: sample.keypoints,
|
|
conf: sample.conf,
|
|
timestamp: sample.timestamp,
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Deterministic shuffle
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function shuffleArray(arr, seed) {
|
|
const result = [...arr];
|
|
let s = seed;
|
|
for (let i = result.length - 1; i > 0; i--) {
|
|
s ^= s << 13; s ^= s >> 17; s ^= s << 5;
|
|
const j = (s >>> 0) % (i + 1);
|
|
[result[i], result[j]] = [result[j], result[i]];
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// WiFlow Supervised Model — simplified TCN + linear decoder
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* 1D causal dilated convolution layer.
|
|
* Weight shape: [outCh, inCh, kernel] stored as flat Float32Array.
|
|
* Input/output layout: [channels, T].
|
|
*/
|
|
class CausalConv1d {
|
|
constructor(inCh, outCh, kernel, dilation, rng) {
|
|
this.inCh = inCh;
|
|
this.outCh = outCh;
|
|
this.kernel = kernel;
|
|
this.dilation = dilation || 1;
|
|
|
|
// Kaiming init
|
|
this.weight = initKaiming(inCh * kernel, outCh, rng);
|
|
this.bias = new Float32Array(outCh);
|
|
|
|
// Momentum buffers for SGD
|
|
this.weightMom = new Float32Array(this.weight.length);
|
|
this.biasMom = new Float32Array(outCh);
|
|
}
|
|
|
|
numParams() {
|
|
return this.weight.length + this.bias.length;
|
|
}
|
|
|
|
/**
|
|
* Forward: [inCh, T] -> [outCh, T] with causal (left) padding.
|
|
*/
|
|
forward(input, T) {
|
|
const effectiveK = this.kernel + (this.kernel - 1) * (this.dilation - 1);
|
|
const padLeft = effectiveK - 1;
|
|
const T_padded = T + padLeft;
|
|
|
|
// Pad input
|
|
const padded = new Float32Array(this.inCh * T_padded);
|
|
for (let c = 0; c < this.inCh; c++) {
|
|
for (let t = 0; t < T; t++) {
|
|
padded[c * T_padded + (t + padLeft)] = input[c * T + t];
|
|
}
|
|
}
|
|
|
|
// Convolve
|
|
const output = new Float32Array(this.outCh * T);
|
|
for (let oc = 0; oc < this.outCh; oc++) {
|
|
for (let t = 0; t < T; t++) {
|
|
let sum = this.bias[oc];
|
|
for (let ic = 0; ic < this.inCh; ic++) {
|
|
for (let k = 0; k < this.kernel; k++) {
|
|
const tIdx = t + padLeft - k * this.dilation;
|
|
if (tIdx >= 0 && tIdx < T_padded) {
|
|
const wIdx = oc * (this.inCh * this.kernel) + ic * this.kernel + k;
|
|
sum += this.weight[wIdx] * padded[ic * T_padded + tIdx];
|
|
}
|
|
}
|
|
}
|
|
output[oc * T + t] = sum;
|
|
}
|
|
}
|
|
return output;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Batch normalization for 1D temporal data [channels, T].
|
|
* Uses running mean/var for inference; batch stats for training.
|
|
*/
|
|
class BatchNorm1d {
|
|
constructor(channels) {
|
|
this.channels = channels;
|
|
this.gamma = new Float32Array(channels).fill(1.0);
|
|
this.beta = new Float32Array(channels);
|
|
this.runMean = new Float32Array(channels);
|
|
this.runVar = new Float32Array(channels).fill(1.0);
|
|
this.momentum = 0.1;
|
|
this.eps = 1e-5;
|
|
|
|
// Momentum buffers
|
|
this.gammaMom = new Float32Array(channels);
|
|
this.betaMom = new Float32Array(channels);
|
|
}
|
|
|
|
numParams() {
|
|
return this.channels * 2;
|
|
}
|
|
|
|
/**
|
|
* Forward: [channels, T] -> [channels, T], updates running stats.
|
|
*/
|
|
forward(input, T) {
|
|
const output = new Float32Array(input.length);
|
|
for (let c = 0; c < this.channels; c++) {
|
|
// Compute channel mean and var over T
|
|
let mean = 0, varAcc = 0;
|
|
for (let t = 0; t < T; t++) mean += input[c * T + t];
|
|
mean /= T;
|
|
for (let t = 0; t < T; t++) varAcc += (input[c * T + t] - mean) ** 2;
|
|
varAcc /= T;
|
|
|
|
// Update running stats
|
|
this.runMean[c] = (1 - this.momentum) * this.runMean[c] + this.momentum * mean;
|
|
this.runVar[c] = (1 - this.momentum) * this.runVar[c] + this.momentum * varAcc;
|
|
|
|
// Normalize
|
|
const invStd = 1.0 / Math.sqrt(varAcc + this.eps);
|
|
for (let t = 0; t < T; t++) {
|
|
output[c * T + t] = this.gamma[c] * (input[c * T + t] - mean) * invStd + this.beta[c];
|
|
}
|
|
}
|
|
return output;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* TCN block: Conv1d (causal, dilated) -> BN -> ReLU -> Conv1d -> BN + residual -> ReLU
|
|
*/
|
|
class TCNBlock {
|
|
constructor(inCh, outCh, kernel, dilation, rng) {
|
|
this.conv1 = new CausalConv1d(inCh, outCh, kernel, dilation, rng);
|
|
this.bn1 = new BatchNorm1d(outCh);
|
|
this.conv2 = new CausalConv1d(outCh, outCh, kernel, dilation, rng);
|
|
this.bn2 = new BatchNorm1d(outCh);
|
|
|
|
// Residual projection if dimensions differ
|
|
this.hasResProj = (inCh !== outCh);
|
|
if (this.hasResProj) {
|
|
this.resConv = new CausalConv1d(inCh, outCh, 1, 1, rng);
|
|
}
|
|
}
|
|
|
|
numParams() {
|
|
let p = this.conv1.numParams() + this.bn1.numParams() +
|
|
this.conv2.numParams() + this.bn2.numParams();
|
|
if (this.hasResProj) p += this.resConv.numParams();
|
|
return p;
|
|
}
|
|
|
|
forward(input, T) {
|
|
// Path 1: conv -> bn -> relu -> conv -> bn
|
|
let x = this.conv1.forward(input, T);
|
|
x = this.bn1.forward(x, T);
|
|
relu(x);
|
|
x = this.conv2.forward(x, T);
|
|
x = this.bn2.forward(x, T);
|
|
|
|
// Residual
|
|
const res = this.hasResProj ? this.resConv.forward(input, T) : input;
|
|
for (let i = 0; i < x.length; i++) x[i] += res[i];
|
|
relu(x);
|
|
return x;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Linear layer: [inDim] -> [outDim]
|
|
*/
|
|
class Linear {
|
|
constructor(inDim, outDim, rng) {
|
|
this.inDim = inDim;
|
|
this.outDim = outDim;
|
|
this.weight = initXavier(inDim, outDim, rng);
|
|
this.bias = new Float32Array(outDim);
|
|
|
|
// Momentum buffers
|
|
this.weightMom = new Float32Array(this.weight.length);
|
|
this.biasMom = new Float32Array(outDim);
|
|
}
|
|
|
|
numParams() {
|
|
return this.weight.length + this.bias.length;
|
|
}
|
|
|
|
forward(input) {
|
|
const output = new Float32Array(this.outDim);
|
|
for (let j = 0; j < this.outDim; j++) {
|
|
let sum = this.bias[j];
|
|
for (let i = 0; i < this.inDim; i++) {
|
|
sum += input[i] * this.weight[i * this.outDim + j];
|
|
}
|
|
output[j] = sum;
|
|
}
|
|
return output;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* WiFlow Supervised Model.
|
|
*
|
|
* TCN Stage: 4 dilated causal conv blocks (dilation 1,2,4,8), kernel 7
|
|
* input_dim -> 256 -> 192 -> 128
|
|
* Flatten + Linear: [128 * 20] -> 2048 -> [17 * 2]
|
|
* Sigmoid to [0, 1]
|
|
*/
|
|
class WiFlowSupervisedModel {
|
|
constructor(inputDim, timeSteps, numKeypoints, seed, scale) {
|
|
this.inputDim = inputDim;
|
|
this.timeSteps = timeSteps;
|
|
this.numKeypoints = numKeypoints || 17;
|
|
this.outDim = this.numKeypoints * 2;
|
|
this.scale = scale || SCALE;
|
|
|
|
const rng = createRng(seed || 42);
|
|
const ch = this.scale.tcnChannels;
|
|
const k = this.scale.kernel;
|
|
|
|
// TCN blocks: inputDim -> ch[0] -> ch[1] -> ch[2] -> ch[3]
|
|
this.tcnBlocks = [];
|
|
let prevCh = inputDim;
|
|
const dilations = [1, 2, 4, 8];
|
|
const nBlocks = Math.min(this.scale.tcnBlocks, ch.length);
|
|
for (let i = 0; i < nBlocks; i++) {
|
|
this.tcnBlocks.push(new TCNBlock(prevCh, ch[i], k, dilations[i], rng));
|
|
prevCh = ch[i];
|
|
}
|
|
|
|
// Flatten: lastCh * timeSteps -> hidden -> 34
|
|
const flatDim = prevCh * timeSteps;
|
|
const hiddenDim = this.scale.hiddenDim;
|
|
this.fc1 = new Linear(flatDim, hiddenDim, rng);
|
|
this.fc2 = new Linear(hiddenDim, this.outDim, rng);
|
|
|
|
this._totalParams = null;
|
|
}
|
|
|
|
totalParams() {
|
|
if (this._totalParams === null) {
|
|
this._totalParams = this.fc1.numParams() + this.fc2.numParams();
|
|
for (const b of this.tcnBlocks) this._totalParams += b.numParams();
|
|
}
|
|
return this._totalParams;
|
|
}
|
|
|
|
/**
|
|
* Forward pass.
|
|
* @param {Float32Array} csi - [inputDim * timeSteps] flat
|
|
* @returns {Float32Array} keypoints [numKeypoints * 2] in [0, 1]
|
|
*/
|
|
forward(csi) {
|
|
const T = this.timeSteps;
|
|
|
|
// TCN stages (dynamic block count based on scale)
|
|
let x = csi;
|
|
for (const block of this.tcnBlocks) {
|
|
x = block.forward(x, T);
|
|
}
|
|
|
|
// FC layers with ReLU
|
|
let h = this.fc1.forward(x);
|
|
relu(h);
|
|
let out = this.fc2.forward(h);
|
|
|
|
// Sigmoid to [0, 1]
|
|
for (let i = 0; i < out.length; i++) {
|
|
out[i] = sigmoid(out[i]);
|
|
}
|
|
|
|
return out;
|
|
}
|
|
|
|
/**
|
|
* Encode CSI to embedding (for contrastive phase).
|
|
* Returns the fc1 hidden layer (2048-dim).
|
|
*/
|
|
encode(csi) {
|
|
const T = this.timeSteps;
|
|
let x = csi;
|
|
for (const block of this.tcnBlocks) {
|
|
x = block.forward(x, T);
|
|
}
|
|
|
|
let h = this.fc1.forward(x);
|
|
relu(h);
|
|
|
|
// L2 normalize for contrastive
|
|
let norm = 0;
|
|
for (let i = 0; i < h.length; i++) norm += h[i] * h[i];
|
|
norm = Math.sqrt(norm) || 1;
|
|
for (let i = 0; i < h.length; i++) h[i] /= norm;
|
|
|
|
return h;
|
|
}
|
|
|
|
/**
|
|
* Collect all weight arrays for gradient updates.
|
|
* Returns array of { weight, mom, name } objects.
|
|
*/
|
|
collectParams() {
|
|
const params = [];
|
|
const addConv = (conv, prefix) => {
|
|
params.push({ weight: conv.weight, mom: conv.weightMom, name: `${prefix}.weight` });
|
|
params.push({ weight: conv.bias, mom: conv.biasMom, name: `${prefix}.bias` });
|
|
};
|
|
const addBN = (bn, prefix) => {
|
|
params.push({ weight: bn.gamma, mom: bn.gammaMom, name: `${prefix}.gamma` });
|
|
params.push({ weight: bn.beta, mom: bn.betaMom, name: `${prefix}.beta` });
|
|
};
|
|
const addTCN = (tcn, prefix) => {
|
|
addConv(tcn.conv1, `${prefix}.conv1`);
|
|
addBN(tcn.bn1, `${prefix}.bn1`);
|
|
addConv(tcn.conv2, `${prefix}.conv2`);
|
|
addBN(tcn.bn2, `${prefix}.bn2`);
|
|
if (tcn.hasResProj) addConv(tcn.resConv, `${prefix}.res`);
|
|
};
|
|
const addLinear = (linear, prefix) => {
|
|
params.push({ weight: linear.weight, mom: linear.weightMom, name: `${prefix}.weight` });
|
|
params.push({ weight: linear.bias, mom: linear.biasMom, name: `${prefix}.bias` });
|
|
};
|
|
|
|
for (let i = 0; i < this.tcnBlocks.length; i++) {
|
|
addTCN(this.tcnBlocks[i], `tcn${i}`);
|
|
}
|
|
addLinear(this.fc1, 'fc1');
|
|
addLinear(this.fc2, 'fc2');
|
|
|
|
return params;
|
|
}
|
|
|
|
/**
|
|
* Get all weights as a flat Float32Array (for export).
|
|
*/
|
|
getAllWeights() {
|
|
const params = this.collectParams();
|
|
let totalLen = 0;
|
|
for (const p of params) totalLen += p.weight.length;
|
|
const flat = new Float32Array(totalLen);
|
|
let offset = 0;
|
|
for (const p of params) {
|
|
flat.set(p.weight, offset);
|
|
offset += p.weight.length;
|
|
}
|
|
return flat;
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// SGD with momentum + cosine LR decay
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Numerical gradient estimation using finite differences.
|
|
* Computes gradient of lossFn w.r.t. each parameter in paramObj.weight.
|
|
*/
|
|
function computeNumericalGrad(model, sample, lossFn, paramObj, eps) {
|
|
eps = eps || 1e-4;
|
|
const w = paramObj.weight;
|
|
const grad = new Float32Array(w.length);
|
|
|
|
for (let i = 0; i < w.length; i++) {
|
|
const orig = w[i];
|
|
|
|
w[i] = orig + eps;
|
|
const lossPlus = lossFn(model, sample);
|
|
|
|
w[i] = orig - eps;
|
|
const lossMinus = lossFn(model, sample);
|
|
|
|
w[i] = orig;
|
|
grad[i] = (lossPlus - lossMinus) / (2 * eps);
|
|
}
|
|
|
|
return grad;
|
|
}
|
|
|
|
/**
|
|
* Apply SGD with momentum to a single parameter.
|
|
*/
|
|
function sgdStep(paramObj, grad, lr, momentum) {
|
|
const w = paramObj.weight;
|
|
const mom = paramObj.mom;
|
|
for (let i = 0; i < w.length; i++) {
|
|
mom[i] = momentum * mom[i] + grad[i];
|
|
w[i] -= lr * mom[i];
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Cosine annealing learning rate.
|
|
*/
|
|
function cosineDecayLR(baseLR, epoch, totalEpochs) {
|
|
return baseLR * 0.5 * (1 + Math.cos(Math.PI * epoch / totalEpochs));
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Loss functions
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Confidence-weighted SmoothL1 loss for keypoints.
|
|
* L = (1/N) * sum(conf_i * smoothL1(pred_i, gt_i, beta=0.05))
|
|
*/
|
|
function supervisedLoss(predicted, target, conf, beta) {
|
|
beta = beta || 0.05;
|
|
const nKp = conf.length;
|
|
let loss = 0;
|
|
let weightSum = 0;
|
|
|
|
for (let k = 0; k < nKp; k++) {
|
|
const px = predicted[k * 2], py = predicted[k * 2 + 1];
|
|
const tx = target[k * 2], ty = target[k * 2 + 1];
|
|
|
|
const diffX = Math.abs(px - tx);
|
|
const diffY = Math.abs(py - ty);
|
|
|
|
let lx = diffX < beta ? 0.5 * diffX * diffX / beta : diffX - 0.5 * beta;
|
|
let ly = diffY < beta ? 0.5 * diffY * diffY / beta : diffY - 0.5 * beta;
|
|
|
|
loss += conf[k] * (lx + ly);
|
|
weightSum += conf[k];
|
|
}
|
|
|
|
return weightSum > 0 ? loss / weightSum : 0;
|
|
}
|
|
|
|
/**
|
|
* Bone length constraint loss.
|
|
*/
|
|
function boneLoss(predicted) {
|
|
let loss = 0;
|
|
for (let b = 0; b < BONE_CONNECTIONS.length; b++) {
|
|
const [i, j] = BONE_CONNECTIONS[b];
|
|
const prior = BONE_LENGTH_PRIORS[b];
|
|
const dx = predicted[i * 2] - predicted[j * 2];
|
|
const dy = predicted[i * 2 + 1] - predicted[j * 2 + 1];
|
|
const boneLen = Math.sqrt(dx * dx + dy * dy);
|
|
const deviation = boneLen - prior;
|
|
loss += deviation * deviation;
|
|
}
|
|
return loss / BONE_CONNECTIONS.length;
|
|
}
|
|
|
|
/**
|
|
* Temporal consistency loss between consecutive predictions.
|
|
*/
|
|
function temporalLoss(predCurrent, predPrev) {
|
|
if (!predPrev) return 0;
|
|
return smoothL1(predCurrent, predPrev, 0.05);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Evaluation: PCK@threshold
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function pck(predicted, target, threshold) {
|
|
threshold = threshold || 0.2;
|
|
let correct = 0;
|
|
const nKp = Math.min(predicted.length, target.length) / 2;
|
|
for (let k = 0; k < nKp; k++) {
|
|
const dx = predicted[k * 2] - target[k * 2];
|
|
const dy = predicted[k * 2 + 1] - target[k * 2 + 1];
|
|
if (Math.sqrt(dx * dx + dy * dy) < threshold) correct++;
|
|
}
|
|
return correct / nKp;
|
|
}
|
|
|
|
/**
|
|
* Evaluate model on held-out set, return average loss and PCK@20.
|
|
*/
|
|
function evaluate(model, evalSet) {
|
|
let totalLoss = 0;
|
|
let totalPck = 0;
|
|
|
|
for (const sample of evalSet) {
|
|
const pred = model.forward(sample.csi);
|
|
totalLoss += supervisedLoss(pred, sample.keypoints, sample.conf);
|
|
totalPck += pck(pred, sample.keypoints, 0.2);
|
|
}
|
|
|
|
return {
|
|
loss: evalSet.length > 0 ? totalLoss / evalSet.length : 0,
|
|
pck20: evalSet.length > 0 ? totalPck / evalSet.length : 0,
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Stochastic gradient estimation for a mini-batch
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Estimate gradient via forward-mode perturbation for a mini-batch.
|
|
* This uses simultaneous perturbation (SPSA-like) which scales O(1) per
|
|
* parameter rather than O(n) for naive numerical differentiation.
|
|
*/
|
|
function estimateBatchGrad(model, batch, lossFn, paramObj, rng) {
|
|
const eps = 1e-4;
|
|
const w = paramObj.weight;
|
|
const n = w.length;
|
|
const grad = new Float32Array(n);
|
|
|
|
// Use SPSA: perturb all weights simultaneously with random direction
|
|
const delta = new Float32Array(n);
|
|
for (let i = 0; i < n; i++) {
|
|
delta[i] = rng() < 0.5 ? 1 : -1;
|
|
}
|
|
|
|
// Compute loss at w + eps*delta
|
|
for (let i = 0; i < n; i++) w[i] += eps * delta[i];
|
|
let lossPlus = 0;
|
|
for (const sample of batch) lossPlus += lossFn(model, sample);
|
|
lossPlus /= batch.length;
|
|
|
|
// Compute loss at w - eps*delta
|
|
for (let i = 0; i < n; i++) w[i] -= 2 * eps * delta[i];
|
|
let lossMinus = 0;
|
|
for (const sample of batch) lossMinus += lossFn(model, sample);
|
|
lossMinus /= batch.length;
|
|
|
|
// Restore weights
|
|
for (let i = 0; i < n; i++) w[i] += eps * delta[i];
|
|
|
|
// SPSA gradient estimate
|
|
const scale = (lossPlus - lossMinus) / (2 * eps);
|
|
for (let i = 0; i < n; i++) {
|
|
grad[i] = scale / delta[i];
|
|
}
|
|
|
|
return grad;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Main training pipeline
|
|
// ---------------------------------------------------------------------------
|
|
|
|
async function main() {
|
|
const startTime = Date.now();
|
|
console.log('=== WiFlow Supervised Pose Training Pipeline (ADR-079) ===');
|
|
console.log(`Config: totalEpochs=${CONFIG.totalEpochs} batch=${CONFIG.batchSize} lr=${CONFIG.lr}`);
|
|
console.log(` phases: contrastive=${contrastiveEpochs} supervised=${supervisedEpochs} refinement=${refinementEpochs}`);
|
|
console.log(` momentum=${CONFIG.momentum} evalSplit=${CONFIG.evalSplit}`);
|
|
console.log('');
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Step 1: Load paired data
|
|
// -----------------------------------------------------------------------
|
|
console.log('[1/6] Loading paired CSI+keypoint data...');
|
|
const allSamples = loadPairedData(CONFIG.dataPath);
|
|
if (allSamples.length === 0) {
|
|
console.error('No valid paired samples found in data file.');
|
|
process.exit(1);
|
|
}
|
|
|
|
// Auto-detect input dimension
|
|
let inputDim = allSamples[0].csiDim;
|
|
const T = CONFIG.timeSteps;
|
|
console.log(` Loaded ${allSamples.length} paired samples`);
|
|
console.log(` Auto-detected input dim: ${inputDim} (${inputDim === 128 ? 'full CSI subcarriers' : inputDim + '-dim feature vectors'})`);
|
|
console.log(` Time steps: ${T}`);
|
|
|
|
// -----------------------------------------------------------------------
|
|
// O6: Subcarrier selection (ruvector-solver inspired)
|
|
// -----------------------------------------------------------------------
|
|
let selectedSubcarriers = null;
|
|
if (inputDim >= 64) {
|
|
const topK = Math.min(56, Math.floor(inputDim * 0.5)); // 50% reduction like ruvector 114→56
|
|
console.log(` [O6] Selecting top-${topK} subcarriers by variance (ruvector-solver)...`);
|
|
selectedSubcarriers = selectTopSubcarriers(allSamples, inputDim, T, topK);
|
|
const origDim = inputDim;
|
|
// Reduce all samples
|
|
for (let i = 0; i < allSamples.length; i++) {
|
|
allSamples[i] = reduceSubcarriers(allSamples[i], selectedSubcarriers, T);
|
|
}
|
|
inputDim = topK;
|
|
console.log(` [O6] Reduced: ${origDim} → ${inputDim} subcarriers (${((1 - inputDim / origDim) * 100).toFixed(0)}% reduction)`);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// O7: Subcarrier attention weighting (ruvector-attention inspired)
|
|
// -----------------------------------------------------------------------
|
|
console.log(` [O7] Computing subcarrier attention weights (ruvector-attention)...`);
|
|
const subcarrierAttention = computeSubcarrierAttention(allSamples, inputDim, T);
|
|
// Apply attention to all samples
|
|
for (let i = 0; i < allSamples.length; i++) {
|
|
allSamples[i].csi = applySubcarrierAttention(allSamples[i].csi, subcarrierAttention, inputDim, T);
|
|
}
|
|
const topAttnIdx = Array.from({ length: inputDim }, (_, i) => i)
|
|
.sort((a, b) => subcarrierAttention[b] - subcarrierAttention[a])
|
|
.slice(0, 5);
|
|
console.log(` [O7] Top-5 attention subcarriers: [${topAttnIdx.join(', ')}]`);
|
|
|
|
// -----------------------------------------------------------------------
|
|
// O8: DynamicMinCut person separation (ruvector-mincut inspired)
|
|
// -----------------------------------------------------------------------
|
|
if (inputDim >= 16) {
|
|
console.log(` [O8] Running Stoer-Wagner min-cut for person separation (ruvector-mincut)...`);
|
|
const mcSamples = allSamples.slice(0, Math.min(50, allSamples.length)); // subsample for speed
|
|
const mcResult = minCutPersonSeparation(mcSamples, inputDim, T);
|
|
const g0 = mcResult.partition.filter(v => v === 0).length;
|
|
const g1 = mcResult.partition.filter(v => v === 1).length;
|
|
console.log(` [O8] Min-cut value: ${mcResult.cutValue.toFixed(4)} — partition: [${g0}, ${g1}] subcarriers`);
|
|
console.log(` [O8] Person-separable subcarrier groups identified for multi-person training`);
|
|
}
|
|
|
|
// Train/eval split
|
|
const shuffled = shuffleArray(allSamples, 42);
|
|
const splitIdx = Math.floor(shuffled.length * (1 - CONFIG.evalSplit));
|
|
const trainSet = shuffled.slice(0, splitIdx);
|
|
const evalSet = shuffled.slice(splitIdx);
|
|
console.log(` Train: ${trainSet.length} Eval: ${evalSet.length}`);
|
|
console.log('');
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Step 2: Initialize model
|
|
// -----------------------------------------------------------------------
|
|
console.log('[2/6] Initializing WiFlow supervised model...');
|
|
const model = new WiFlowSupervisedModel(inputDim, T, CONFIG.numKeypoints, 42, SCALE);
|
|
const ch = SCALE.tcnChannels.slice(0, SCALE.tcnBlocks);
|
|
const lastCh = ch[ch.length - 1];
|
|
console.log(` Scale: ${scaleKey}`);
|
|
console.log(` Parameters: ${model.totalParams().toLocaleString()}`);
|
|
console.log(` Architecture: TCN(${inputDim}->${ch.join('->')}, k=${SCALE.kernel}, d=[1,2,4,8]) -> FC(${lastCh * T}->${SCALE.hiddenDim}->34)`);
|
|
console.log('');
|
|
|
|
const trainingLog = {
|
|
config: { ...CONFIG, inputDim, contrastiveEpochs, supervisedEpochs, refinementEpochs },
|
|
phases: [],
|
|
};
|
|
|
|
const allParams = model.collectParams();
|
|
const rng = createRng(123);
|
|
let globalEpoch = 0;
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Phase 1: Contrastive pretraining
|
|
// -----------------------------------------------------------------------
|
|
if (!CONFIG.skipContrastive && contrastiveEpochs > 0) {
|
|
console.log(`[3/6] Phase 1: Contrastive pretraining (${contrastiveEpochs} epochs)...`);
|
|
|
|
const contrastiveLog = { phase: 'contrastive', epochs: [] };
|
|
const trainer = new ContrastiveTrainer({
|
|
margin: 0.3,
|
|
temperature: 0.07,
|
|
});
|
|
|
|
for (let epoch = 0; epoch < contrastiveEpochs; epoch++) {
|
|
const lr = cosineDecayLR(CONFIG.lr * 10, epoch, contrastiveEpochs); // Higher LR for contrastive
|
|
const shuffledTrain = shuffleArray(trainSet, epoch * 7 + 1);
|
|
|
|
let epochLoss = 0;
|
|
let nBatches = 0;
|
|
|
|
for (let b = 0; b < shuffledTrain.length - 2; b += CONFIG.batchSize) {
|
|
const batchEnd = Math.min(b + CONFIG.batchSize, shuffledTrain.length - 2);
|
|
let batchLoss = 0;
|
|
let nTriplets = 0;
|
|
|
|
// Create temporal triplets: anchor=frame[i], positive=frame[i+1], negative=frame[j] (far)
|
|
for (let i = b; i < batchEnd; i++) {
|
|
const anchorEmb = Array.from(model.encode(shuffledTrain[i].csi));
|
|
const positiveEmb = Array.from(model.encode(shuffledTrain[i + 1].csi));
|
|
// Negative: pick a distant sample
|
|
const negIdx = (i + Math.floor(shuffledTrain.length / 2)) % shuffledTrain.length;
|
|
const negativeEmb = Array.from(model.encode(shuffledTrain[negIdx].csi));
|
|
|
|
trainer.addTriplet(
|
|
`anchor-${i}`, anchorEmb,
|
|
`pos-${i}`, positiveEmb,
|
|
`neg-${i}`, negativeEmb,
|
|
);
|
|
|
|
const sim_pos = cosineSimilarity(anchorEmb, positiveEmb);
|
|
const sim_neg = cosineSimilarity(anchorEmb, negativeEmb);
|
|
batchLoss += Math.max(0, 0.3 - sim_pos + sim_neg);
|
|
nTriplets++;
|
|
}
|
|
|
|
if (nTriplets > 0) batchLoss /= nTriplets;
|
|
|
|
// SPSA gradient update on all params
|
|
for (const p of allParams) {
|
|
const lossFn = (m, s) => {
|
|
const emb = m.encode(s.csi);
|
|
// Simple self-consistency loss
|
|
let norm = 0;
|
|
for (let i = 0; i < emb.length; i++) norm += emb[i] * emb[i];
|
|
return 1.0 - norm; // push toward unit norm
|
|
};
|
|
|
|
const batch = shuffledTrain.slice(b, batchEnd);
|
|
const grad = multiSpsaGrad(model, batch, lossFn, p, rng, SCALE.spsaK);
|
|
sgdStep(p, grad, lr, CONFIG.momentum);
|
|
}
|
|
|
|
epochLoss += batchLoss;
|
|
nBatches++;
|
|
}
|
|
|
|
epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
|
|
const evalResult = evaluate(model, evalSet);
|
|
|
|
contrastiveLog.epochs.push({
|
|
epoch: globalEpoch,
|
|
loss: epochLoss,
|
|
evalLoss: evalResult.loss,
|
|
pck20: evalResult.pck20,
|
|
lr,
|
|
});
|
|
|
|
if ((epoch + 1) % 10 === 0 || epoch === 0) {
|
|
console.log(` [contrastive] epoch ${epoch + 1}/${contrastiveEpochs} loss=${epochLoss.toFixed(6)} eval_loss=${evalResult.loss.toFixed(6)} PCK@20=${(evalResult.pck20 * 100).toFixed(1)}% lr=${lr.toExponential(2)}`);
|
|
}
|
|
globalEpoch++;
|
|
}
|
|
|
|
trainingLog.phases.push(contrastiveLog);
|
|
console.log('');
|
|
} else {
|
|
console.log('[3/6] Phase 1: Contrastive pretraining SKIPPED');
|
|
console.log('');
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Phase 2: Supervised training with curriculum (O1)
|
|
// -----------------------------------------------------------------------
|
|
console.log(`[4/6] Phase 2: Supervised keypoint regression (${supervisedEpochs} epochs, 4-stage curriculum)...`);
|
|
|
|
const supervisedLog = { phase: 'supervised', epochs: [] };
|
|
const epochsPerStage = Math.floor(supervisedEpochs / CONFIG.curriculumStages.length);
|
|
|
|
for (let epoch = 0; epoch < supervisedEpochs; epoch++) {
|
|
// Determine curriculum stage
|
|
const stageIdx = Math.min(
|
|
Math.floor(epoch / epochsPerStage),
|
|
CONFIG.curriculumStages.length - 1
|
|
);
|
|
const confThreshold = CONFIG.curriculumStages[stageIdx];
|
|
const useAugmentation = (stageIdx === CONFIG.curriculumStages.length - 1);
|
|
|
|
const lr = cosineDecayLR(CONFIG.lr, epoch, supervisedEpochs);
|
|
|
|
// Filter training samples by confidence threshold
|
|
let trainSubset;
|
|
if (confThreshold > 0) {
|
|
trainSubset = trainSet.filter(s => {
|
|
let meanConf = 0;
|
|
for (let i = 0; i < s.conf.length; i++) meanConf += s.conf[i];
|
|
meanConf /= s.conf.length;
|
|
return meanConf >= confThreshold;
|
|
});
|
|
} else {
|
|
trainSubset = trainSet;
|
|
}
|
|
|
|
// Apply augmentation in final stage
|
|
if (useAugmentation) {
|
|
const augmented = [];
|
|
for (const s of trainSubset) {
|
|
augmented.push(s);
|
|
augmented.push(augmentSample(s, createRng(epoch * 1000 + augmented.length), T));
|
|
}
|
|
trainSubset = augmented;
|
|
}
|
|
|
|
if (trainSubset.length === 0) {
|
|
// Skip if no samples pass threshold
|
|
globalEpoch++;
|
|
continue;
|
|
}
|
|
|
|
const shuffledTrain = shuffleArray(trainSubset, epoch * 13 + 3);
|
|
|
|
let epochLoss = 0;
|
|
let nBatches = 0;
|
|
|
|
for (let b = 0; b < shuffledTrain.length; b += CONFIG.batchSize) {
|
|
const batchEnd = Math.min(b + CONFIG.batchSize, shuffledTrain.length);
|
|
const batch = shuffledTrain.slice(b, batchEnd);
|
|
|
|
// Compute batch loss
|
|
const lossFn = (m, s) => {
|
|
const pred = m.forward(s.csi);
|
|
return supervisedLoss(pred, s.keypoints, s.conf);
|
|
};
|
|
|
|
let batchLoss = 0;
|
|
for (const s of batch) batchLoss += lossFn(model, s);
|
|
batchLoss /= batch.length;
|
|
|
|
// SPSA gradient update
|
|
for (const p of allParams) {
|
|
const grad = estimateBatchGrad(model, batch, lossFn, p, rng);
|
|
sgdStep(p, grad, lr, CONFIG.momentum);
|
|
}
|
|
|
|
epochLoss += batchLoss;
|
|
nBatches++;
|
|
}
|
|
|
|
epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
|
|
const evalResult = evaluate(model, evalSet);
|
|
|
|
supervisedLog.epochs.push({
|
|
epoch: globalEpoch,
|
|
stage: stageIdx + 1,
|
|
confThreshold,
|
|
loss: epochLoss,
|
|
evalLoss: evalResult.loss,
|
|
pck20: evalResult.pck20,
|
|
lr,
|
|
trainSamples: trainSubset.length,
|
|
});
|
|
|
|
if ((epoch + 1) % 10 === 0 || epoch === 0) {
|
|
console.log(` [supervised] epoch ${epoch + 1}/${supervisedEpochs} stage=${stageIdx + 1}/4 (conf>${confThreshold.toFixed(1)}) loss=${epochLoss.toFixed(6)} eval_loss=${evalResult.loss.toFixed(6)} PCK@20=${(evalResult.pck20 * 100).toFixed(1)}% lr=${lr.toExponential(2)} samples=${trainSubset.length}`);
|
|
}
|
|
globalEpoch++;
|
|
}
|
|
|
|
trainingLog.phases.push(supervisedLog);
|
|
console.log('');
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Phase 3: Refinement with bone + temporal constraints
|
|
// -----------------------------------------------------------------------
|
|
console.log(`[5/6] Phase 3: Refinement with bone + temporal constraints (${refinementEpochs} epochs)...`);
|
|
|
|
const refinementLog = { phase: 'refinement', epochs: [] };
|
|
|
|
for (let epoch = 0; epoch < refinementEpochs; epoch++) {
|
|
const lr = cosineDecayLR(CONFIG.lr * 0.5, epoch, refinementEpochs); // Lower LR
|
|
const shuffledTrain = shuffleArray(trainSet, epoch * 17 + 7);
|
|
|
|
// Apply augmentation
|
|
const augmented = [];
|
|
for (const s of shuffledTrain) {
|
|
augmented.push(s);
|
|
augmented.push(augmentSample(s, createRng(epoch * 2000 + augmented.length), T));
|
|
}
|
|
|
|
let epochLoss = 0;
|
|
let epochBone = 0;
|
|
let epochTemporal = 0;
|
|
let nBatches = 0;
|
|
|
|
for (let b = 0; b < augmented.length; b += CONFIG.batchSize) {
|
|
const batchEnd = Math.min(b + CONFIG.batchSize, augmented.length);
|
|
const batch = augmented.slice(b, batchEnd);
|
|
|
|
// Combined loss function
|
|
const lossFn = (m, s, prevPred) => {
|
|
const pred = m.forward(s.csi);
|
|
const lSup = supervisedLoss(pred, s.keypoints, s.conf);
|
|
const lBone = boneLoss(pred);
|
|
const lTemp = prevPred ? temporalLoss(pred, prevPred) : 0;
|
|
return lSup + CONFIG.boneWeight * lBone + CONFIG.temporalWeight * lTemp;
|
|
};
|
|
|
|
// Compute batch loss with temporal tracking
|
|
let batchLoss = 0;
|
|
let batchBone = 0;
|
|
let batchTemporal = 0;
|
|
let prevPred = null;
|
|
for (const s of batch) {
|
|
const pred = model.forward(s.csi);
|
|
const lSup = supervisedLoss(pred, s.keypoints, s.conf);
|
|
const lBone = boneLoss(pred);
|
|
const lTemp = prevPred ? temporalLoss(pred, prevPred) : 0;
|
|
batchLoss += lSup + CONFIG.boneWeight * lBone + CONFIG.temporalWeight * lTemp;
|
|
batchBone += lBone;
|
|
batchTemporal += lTemp;
|
|
prevPred = pred;
|
|
}
|
|
batchLoss /= batch.length;
|
|
batchBone /= batch.length;
|
|
batchTemporal /= batch.length;
|
|
|
|
// SPSA gradient update with combined loss
|
|
const combinedLossFn = (m, s) => {
|
|
const pred = m.forward(s.csi);
|
|
return supervisedLoss(pred, s.keypoints, s.conf) +
|
|
CONFIG.boneWeight * boneLoss(pred);
|
|
};
|
|
|
|
for (const p of allParams) {
|
|
const grad = estimateBatchGrad(model, batch, combinedLossFn, p, rng);
|
|
sgdStep(p, grad, lr, CONFIG.momentum);
|
|
}
|
|
|
|
epochLoss += batchLoss;
|
|
epochBone += batchBone;
|
|
epochTemporal += batchTemporal;
|
|
nBatches++;
|
|
}
|
|
|
|
epochLoss = nBatches > 0 ? epochLoss / nBatches : 0;
|
|
epochBone = nBatches > 0 ? epochBone / nBatches : 0;
|
|
epochTemporal = nBatches > 0 ? epochTemporal / nBatches : 0;
|
|
const evalResult = evaluate(model, evalSet);
|
|
|
|
refinementLog.epochs.push({
|
|
epoch: globalEpoch,
|
|
loss: epochLoss,
|
|
boneLoss: epochBone,
|
|
temporalLoss: epochTemporal,
|
|
evalLoss: evalResult.loss,
|
|
pck20: evalResult.pck20,
|
|
lr,
|
|
});
|
|
|
|
if ((epoch + 1) % 10 === 0 || epoch === 0) {
|
|
console.log(` [refinement] epoch ${epoch + 1}/${refinementEpochs} loss=${epochLoss.toFixed(6)} bone=${epochBone.toFixed(6)} temporal=${epochTemporal.toFixed(6)} eval_loss=${evalResult.loss.toFixed(6)} PCK@20=${(evalResult.pck20 * 100).toFixed(1)}% lr=${lr.toExponential(2)}`);
|
|
}
|
|
globalEpoch++;
|
|
}
|
|
|
|
trainingLog.phases.push(refinementLog);
|
|
console.log('');
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Step 6: Export
|
|
// -----------------------------------------------------------------------
|
|
console.log('[6/6] Exporting model and results...');
|
|
|
|
fs.mkdirSync(CONFIG.outputDir, { recursive: true });
|
|
|
|
// Export model weights as JSON
|
|
const weights = model.getAllWeights();
|
|
const modelExport = {
|
|
format: 'wiflow-supervised-v1',
|
|
adr: 'ADR-079',
|
|
architecture: {
|
|
inputDim,
|
|
timeSteps: T,
|
|
numKeypoints: CONFIG.numKeypoints,
|
|
tcnChannels: [inputDim, 256, 256, 192, 128],
|
|
tcnKernel: 7,
|
|
tcnDilations: [1, 2, 4, 8],
|
|
fcDims: [128 * T, 2048, CONFIG.numKeypoints * 2],
|
|
},
|
|
totalParams: model.totalParams(),
|
|
weightsBase64: Buffer.from(weights.buffer).toString('base64'),
|
|
trainingSamples: trainSet.length,
|
|
evalSamples: evalSet.length,
|
|
createdAt: new Date().toISOString(),
|
|
};
|
|
|
|
const modelPath = path.join(CONFIG.outputDir, 'wiflow-v1.json');
|
|
fs.writeFileSync(modelPath, JSON.stringify(modelExport, null, 2));
|
|
console.log(` Model weights: ${modelPath} (${(fs.statSync(modelPath).size / 1024).toFixed(0)} KB)`);
|
|
|
|
// Export training log
|
|
const logPath = path.join(CONFIG.outputDir, 'training-log.json');
|
|
fs.writeFileSync(logPath, JSON.stringify(trainingLog, null, 2));
|
|
console.log(` Training log: ${logPath}`);
|
|
|
|
// Export held-out predictions
|
|
const evalPath = path.join(CONFIG.outputDir, 'eval-holdout.jsonl');
|
|
const evalLines = [];
|
|
for (const sample of evalSet) {
|
|
const pred = model.forward(sample.csi);
|
|
const pckScore = pck(pred, sample.keypoints, 0.2);
|
|
evalLines.push(JSON.stringify({
|
|
timestamp: sample.timestamp,
|
|
predicted: Array.from(pred),
|
|
groundTruth: Array.from(sample.keypoints),
|
|
conf: Array.from(sample.conf),
|
|
pck20: pckScore,
|
|
}));
|
|
}
|
|
fs.writeFileSync(evalPath, evalLines.join('\n') + '\n');
|
|
console.log(` Eval holdout: ${evalPath} (${evalSet.length} samples)`);
|
|
|
|
// Final evaluation summary
|
|
const finalEval = evaluate(model, evalSet);
|
|
const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
|
|
|
|
console.log('');
|
|
console.log('=== Training Complete ===');
|
|
console.log(` Total epochs: ${globalEpoch}`);
|
|
console.log(` Final eval loss: ${finalEval.loss.toFixed(6)}`);
|
|
console.log(` Final PCK@20: ${(finalEval.pck20 * 100).toFixed(1)}%`);
|
|
console.log(` Total parameters: ${model.totalParams().toLocaleString()}`);
|
|
console.log(` Elapsed: ${elapsed}s`);
|
|
}
|
|
|
|
main().catch(err => {
|
|
console.error('Training failed:', err);
|
|
process.exit(1);
|
|
});
|