feat(playground/nisps): warm-start weight transfer on MLP resize (meml-8jc)
Add WasmIML.extractWeights() and WasmIML.createWithWarmStart() to preserve learned joystick mappings across output-count changes; resizeMLP() now transfers hidden-layer weights and shared output nodes instead of cold-starting.
This commit is contained in:
parent
2d2cabe159
commit
c69b46f8cc
2 changed files with 80 additions and 8 deletions
|
|
@ -724,35 +724,45 @@ function outputCountForMode(mode) {
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Recreate IML instances with a new output count.
|
* Recreate IML instances with a new output count.
|
||||||
* Resets all weights and training data.
|
* Joystick IML uses warm-start weight transfer to preserve learned mappings.
|
||||||
|
* Training examples are always cleared (dataset is JS-side and output-count-specific).
|
||||||
*/
|
*/
|
||||||
async function resizeMLP(newOutputCount) {
|
async function resizeMLP(newOutputCount) {
|
||||||
if (newOutputCount === N_OUTPUTS) return;
|
if (newOutputCount === N_OUTPUTS) return;
|
||||||
N_OUTPUTS = newOutputCount;
|
N_OUTPUTS = newOutputCount;
|
||||||
|
|
||||||
|
// Extract joystick weights before destroying (for warm-start transfer)
|
||||||
|
const joySnapshot = imlJoy ? imlJoy.extractWeights() : null;
|
||||||
|
|
||||||
// Destroy old IML instances (free WASM memory)
|
// Destroy old IML instances (free WASM memory)
|
||||||
if (imlJoy) imlJoy.destroy();
|
if (imlJoy) imlJoy.destroy();
|
||||||
if (imlHand) imlHand.destroy();
|
if (imlHand) imlHand.destroy();
|
||||||
|
|
||||||
|
// Joystick IML: warm-start from previous weights when possible
|
||||||
|
if (joySnapshot) {
|
||||||
|
imlJoy = await WasmIML.createWithWarmStart(joySnapshot, N_OUTPUTS, 1000, 1.0, 0.00001);
|
||||||
|
} else {
|
||||||
imlJoy = await WasmIML.create(N_JOY_INPUTS, N_OUTPUTS, [32, 48, 64], 1000, 1.0, 0.00001);
|
imlJoy = await WasmIML.create(N_JOY_INPUTS, N_OUTPUTS, [32, 48, 64], 1000, 1.0, 0.00001);
|
||||||
|
imlJoy.randomiseWeights(spreadLevel);
|
||||||
|
}
|
||||||
imlJoy.setLogger(msg => console.log('[NISPS:joy]', msg));
|
imlJoy.setLogger(msg => console.log('[NISPS:joy]', msg));
|
||||||
|
|
||||||
|
// Hand IML: fresh init (warm-start for 14-input networks is a future concern)
|
||||||
imlHand = await WasmIML.create(N_HAND_INPUTS, N_OUTPUTS, [48, 48, 64], 1000, 1.0, 0.00001);
|
imlHand = await WasmIML.create(N_HAND_INPUTS, N_OUTPUTS, [48, 48, 64], 1000, 1.0, 0.00001);
|
||||||
imlHand.setLogger(msg => console.log('[NISPS:hand]', msg));
|
imlHand.setLogger(msg => console.log('[NISPS:hand]', msg));
|
||||||
|
imlHand.randomiseWeights(spreadLevel);
|
||||||
|
|
||||||
iml = (inputMode === 'joystick') ? imlJoy : imlHand;
|
iml = (inputMode === 'joystick') ? imlJoy : imlHand;
|
||||||
|
|
||||||
// Reset dependent state
|
// Reset dependent state
|
||||||
rawParamValues = new Array(N_OUTPUTS).fill(0.5);
|
rawParamValues = new Array(N_OUTPUTS).fill(0.5);
|
||||||
|
|
||||||
// Randomize with current spread
|
|
||||||
imlJoy.randomiseWeights(spreadLevel);
|
|
||||||
imlHand.randomiseWeights(spreadLevel);
|
|
||||||
|
|
||||||
// Re-run inference
|
// Re-run inference
|
||||||
iml.setInput(0, joyX);
|
iml.setInput(0, joyX);
|
||||||
iml.setInput(1, joyY);
|
iml.setInput(1, joyY);
|
||||||
iml.process();
|
iml.process();
|
||||||
|
|
||||||
console.log(`[NISPS] MLP resized to ${N_OUTPUTS} outputs`);
|
console.log(`[NISPS] MLP resized to ${N_OUTPUTS} outputs (joystick IML warm-started)`);
|
||||||
}
|
}
|
||||||
|
|
||||||
// ---- Init ----
|
// ---- Init ----
|
||||||
|
|
@ -1842,7 +1852,7 @@ async function setOutputMode(mode, { skipConfirm = false } = {}) {
|
||||||
|
|
||||||
// Warn about weight reset if resizing (skip during state restore)
|
// Warn about weight reset if resizing (skip during state restore)
|
||||||
if (needsResize && !skipConfirm && iml.dataset.features.length > 0) {
|
if (needsResize && !skipConfirm && iml.dataset.features.length > 0) {
|
||||||
if (!confirm(`Switching to ${mode} mode requires ${targetOutputs} outputs (currently ${N_OUTPUTS}). This will reset the neural network weights and training data. Continue?`)) {
|
if (!confirm(`Switching to ${mode} mode requires ${targetOutputs} outputs (currently ${N_OUTPUTS}). This will reset training examples. Network weights will be partially preserved. Continue?`)) {
|
||||||
syncOutputToggles(outputMode); // revert pill UI
|
syncOutputToggles(outputMode); // revert pill UI
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -405,6 +405,68 @@ export class WasmIML {
|
||||||
this.process();
|
this.process();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ---- Public weight snapshot for warm-start transfer ----
|
||||||
|
/**
|
||||||
|
* Returns a plain object describing the full network weights for later
|
||||||
|
* reinjection via WasmIML.createWithWarmStart().
|
||||||
|
*/
|
||||||
|
extractWeights() {
|
||||||
|
return {
|
||||||
|
layerSizes: [...this.layerSizes], // e.g. [3, 32, 48, 64, 126]
|
||||||
|
weights: this._getFlatWeights(), // plain Array from fromHeapF32
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Async static factory: create a new WasmIML with newOutputCount outputs,
|
||||||
|
* transferring as much of snapshot.weights as possible.
|
||||||
|
* Hidden-layer weights are copied unchanged; output nodes beyond the old
|
||||||
|
* count retain their random-initialised values.
|
||||||
|
*/
|
||||||
|
static async createWithWarmStart(snapshot, newOutputCount, maxIter, learningRate, convergenceThreshold) {
|
||||||
|
const oldLayerSizes = snapshot.layerSizes;
|
||||||
|
const nInputs = oldLayerSizes[0] - 1; // stored with bias (+1), strip it
|
||||||
|
const hiddenLayers = oldLayerSizes.slice(1, -1); // e.g. [32, 48, 64]
|
||||||
|
const oldOutputCount = oldLayerSizes[oldLayerSizes.length - 1];
|
||||||
|
|
||||||
|
// Create fresh instance with new output count
|
||||||
|
const newIml = await WasmIML.create(nInputs, newOutputCount, hiddenLayers, maxIter, learningRate, convergenceThreshold);
|
||||||
|
|
||||||
|
// Calculate prefix weight count (all layers except the final output layer).
|
||||||
|
// Each layer l: n_nodes[l] * (n_nodes[l-1] + 1) weights (inputs + bias).
|
||||||
|
// fullOldSizes: [nInputs+1, ...hiddenLayers, oldOutputCount]
|
||||||
|
let prefixCount = 0;
|
||||||
|
for (let l = 1; l < oldLayerSizes.length - 1; l++) {
|
||||||
|
prefixCount += oldLayerSizes[l] * (oldLayerSizes[l - 1] + 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
const lastHidden = hiddenLayers[hiddenLayers.length - 1]; // e.g. 64
|
||||||
|
const weightsPerOutputNode = lastHidden + 1; // inputs from last hidden + bias
|
||||||
|
|
||||||
|
// Build new weight array, starting from random init
|
||||||
|
const newWeights = newIml._getFlatWeights();
|
||||||
|
const oldWeights = snapshot.weights;
|
||||||
|
|
||||||
|
// Copy hidden layer weights unchanged
|
||||||
|
for (let i = 0; i < prefixCount; i++) {
|
||||||
|
newWeights[i] = oldWeights[i];
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy output layer weights for nodes that existed in the old network
|
||||||
|
const sharedOutputNodes = Math.min(oldOutputCount, newOutputCount);
|
||||||
|
for (let n = 0; n < sharedOutputNodes; n++) {
|
||||||
|
const oldOff = prefixCount + n * weightsPerOutputNode;
|
||||||
|
const newOff = prefixCount + n * weightsPerOutputNode;
|
||||||
|
for (let w = 0; w < weightsPerOutputNode; w++) {
|
||||||
|
newWeights[newOff + w] = oldWeights[oldOff + w];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Nodes beyond old count retain their random init — good for exploration
|
||||||
|
|
||||||
|
newIml._setFlatWeights(newWeights);
|
||||||
|
return newIml;
|
||||||
|
}
|
||||||
|
|
||||||
// ---- Flat weight get/set (for storedWeights save/restore) ----
|
// ---- Flat weight get/set (for storedWeights save/restore) ----
|
||||||
_getFlatWeights() {
|
_getFlatWeights() {
|
||||||
const ptr = this._w.alloc(this._weightCount);
|
const ptr = this._w.alloc(this._weightCount);
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue