mini-jstorch 2.0.3 → 2.0.4

package/README.md

@@ -1,4 +1,4 @@
-## Mini-JSTorch (v2.0.2)
+## Mini-JSTorch (v2.0.4)
 
 ---
 
@@ -9,12 +9,10 @@ This project prioritizes `clarity`, `numerical correctness`, and `accessibility`
 
 ## Changelog
 
-**v2.0.2:**
-- **Fixed critical training bug:** Optimizers (Adam, SGD, AdamW, Lion) now correctly update Linear and Conv2D layer weights
-- **Fixed BatchNorm2d:** Inference mode no longer produces NaN for multi-channel inputs
-- **Fixed ELU activation:** Backward pass now uses correct derivative formula
-- **Fixed saveModel/loadModel:** Now correctly saves and restores all layer types including Conv2D and BatchNorm2d
-- **Fixed BatchNorm2d gradient zeroing:** gradWeight/gradBias now correctly reset between batches
+**v2.0.4:**
+- **Optimized saveModel/loadModel:** Complete rewrite with flat 1D serialization using shape metadata, single-cursor deserialization, and pre-allocated arrays. Reduces JSON size by ~40-50% for large models and improves load performance.
+- **Fixed stateDict/loadStateDict:** Now uses `parameters()` for universal layer compatibility. Previously only saved `W` and `b` properties, silently dropping Conv2D and BatchNorm2d weights.
+
 
 **⚠️ BREAKING CHANGES in v2.0.0:**
 - Tokenizer API: `tokenizeBatch()` → `transform()`, `detokenizeBatch()` → `inverseTransform()`
@@ -74,7 +72,7 @@ In Browser/Website:
     <div id="res"></div>
 
     <script type="module">
-        import { Sequential, Linear, ReLU, MSELoss, Adam, StepLR, Tanh } from 'https://unpkg.com/mini-jstorch@1.8.0/index.js';
+        import { Sequential, Linear, ReLU, MSELoss, Adam, StepLR, Tanh } from 'https://unpkg.com/mini-jstorch@2.0.4/index.js';
         
         async function train() {
             const statusEl = document.getElementById('status');
@@ -187,7 +185,7 @@ In Browser/Website:
 
 ## Node.js
 ```bash
-npm install mini-jstorch
+npm install mini-jstorch@latest
 ```
 Node.js v18+ or any modern browser with ES module support is recommended.
 
@@ -214,13 +212,12 @@ import {
 const model = new Sequential([
   new Linear(2, 8),
   new ReLU(),
-  new Linear(8, 2) // logits output
+  new Linear(8, 2)
 ]);
 
 const X = [
   [0,0], [0,1], [1,0], [1,1]
 ];
-
 const Y = [
   [1,0], [0,1], [0,1], [1,0]
 ];
@@ -234,14 +231,20 @@ for (let epoch = 1; epoch <= 300; epoch++) {
   const grad = lossFn.backward();
   model.backward(grad);
   optimizer.step();
-  
-  // Zero gradients for next iteration
   model.zeroGrad();
 
   if (epoch % 50 === 0) {
     console.log(`Epoch ${epoch}, Loss: ${loss.toFixed(6)}`);
   }
 }
+
+console.log('\nResults:');
+const logits = model.forward(X);
+X.forEach((input, i) => {
+  const pred = logits[i][0] > logits[i][1] ? 0 : 1;
+  const target = Y[i][0] === 1 ? 0 : 1;
+  console.log(`  [${input}] → class ${pred} (target: ${target}) ${pred === target ? 'TRUE' : 'FALSE'}`);
+});
 ```
 `Important:` Do not combine `SoftmaxCrossEntropyLoss` with a `Softmax` layer.
 
@@ -280,7 +283,6 @@ for (let epoch = 1; epoch <= 300; epoch++) {
   optimizer.step();
   model.zeroGrad();
   
-  // Print progress every 50 epochs
   if (epoch % 50 === 0) {
     const probs = logits.map(p => 1 / (1 + Math.exp(-p[0])));
     console.log(`Epoch ${epoch} | Loss: ${loss.toFixed(6)}`);
@@ -292,7 +294,6 @@ for (let epoch = 1; epoch <= 300; epoch++) {
   }
 }
 
-// Final evaluation
 console.log("\nTraining Complete\n");
 model.eval(); 
 
@@ -340,7 +341,7 @@ See the `demo/` directory for runnable examples!
 - `demo/scheduler.js`
 - `demo/xor_classification.js`
 - `demo/linear_regression.js`
-
+- `demo/saveAndLoadModel.js`
 
 ```bash
 node demo/<fileNameInDemo>.js

package/demo/saveAndLoadModel.js

@@ -0,0 +1,45 @@
+// demo/saveAndLoadModel.js
+// Trains a linear regression model (y = 3x + 5), saves it to JSON,
+// loads it back into a fresh model, and verifies predictions match.
+
+import {
+    Sequential, Linear,
+    MSELoss, Adam,
+    saveModel, loadModel
+} from '../src/jstorch.js';
+
+const X = [[1], [2], [3], [4], [5]];
+const y = [[8], [11], [14], [17], [20]];
+
+const model = new Sequential([new Linear(1, 1)]);
+const criterion = new MSELoss();
+const optimizer = new Adam(model.parameters(), 0.1);
+
+// Train
+for (let e = 1; e <= 500; e++) {
+    const pred = model.forward(X);
+    const loss = criterion.forward(pred, y);
+    model.backward(criterion.backward());
+    optimizer.step();
+    model.zeroGrad();
+    if (e % 100 === 0) console.log(`Epoch ${e} | Loss: ${loss.toFixed(6)}`);
+}
+
+// Original predictions
+console.log('\nTrained model predictions:');
+const original = model.forward(X);
+X.forEach((x, i) => console.log(`  x=${x[0]} → ${original[i][0].toFixed(4)} (target: ${y[i][0]})`));
+
+// Save → Load
+const json = saveModel(model);
+const restored = new Sequential([new Linear(1, 1)]);
+loadModel(restored, json);
+
+// Loaded predictions
+console.log('\nLoaded model predictions:');
+const loaded = restored.forward(X);
+X.forEach((x, i) => console.log(`  x=${x[0]} → ${loaded[i][0].toFixed(4)} (target: ${y[i][0]})`));
+
+// Verify
+const ok = original.every((r, i) => Math.abs(r[0] - loaded[i][0]) < 1e-10);
+console.log(`\n${ok ? 'Model saved and restored successfully.' : 'Mismatch after load.'}`);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mini-jstorch",
-  "version": "2.0.3",
+  "version": "2.0.4",
   "type": "module", 
   "description": "A lightweight JavaScript neural network library for learning AI concepts and rapid Frontend experimentation. PyTorch-inspired, zero dependencies, perfect for educational use.",
   "main": "index.js",

package/src/jstorch.js

@@ -703,13 +703,25 @@ export class Sequential {
 	* Get model state dict (weights and biases)
 	*/ 
 	stateDict(){
-		const state = {};
-		this.layers.forEach((layer, idx) => {
-			if (layer.W) state[`layer_${idx}.weight`] = layer.W;
-			if (layer.b) state[`layer_${idx}.bias`] = layer.b;
-		});
-		return state;
-	}
+        const state = {};
+        this.layers.forEach((layer, idx) => {
+            const params = layer.parameters ? layer.parameters() : [];
+            if (params.length === 0) return;
+            
+            params.forEach((p, pIdx) => {
+                // Deep clone parameter data to prevent reference sharing 
+                const paramData = p.param;
+                const is2D = Array.isArray(paramData[0]);
+                
+                if (is2D){
+                    state[`layer_${idx}.param_${pIdx}`] = paramData.map(row => [...row]);
+                } else {
+                    state[`layer_${idx}.param_${pIdx}`] = [...paramData];
+                }
+            });
+        });
+        return state;
+    }
 	
 	step(lr){
 		this.layers.forEach(layer => {
@@ -720,24 +732,70 @@ export class Sequential {
 	}
 	
 	/**
-	* Load state dict 
+	* Load state dict into model.
+    * Mutates parameter in-place to preserve optimizer references.
 	*/
 	loadStateDict(stateDict){
 		this.layers.forEach((layer, idx) => {
-			const weightKey = `layer_${idx}.weight`;
-			const biasKey = `layer_${idx}.bias`;
-			
-			if (layer.W && stateDict[weightKey]){
-				layer.W = stateDict[weightKey];
-				// Invalidate cache 
-				if (layer._InvalidateCache) layer._InvalidateCache();
-			}
-			if (layer.b && stateDict[biasKey]){
-				layer.b = stateDict[biasKey];
-				if (layer._InvalidateCache) layer._InvalidateCache();
-			}
-		});
-		return this;
+            const params = layer.parameters ? layer.parameters() : [];
+            if (params.length === 0) return;
+            
+            params.forEach((p, pIdx) => {
+                const key = `layer_${idx}.param_${pIdx}`;
+                const savedData = stateDict[key];
+                
+                if (!savedData){
+                    console.warn(`[JST WARN]: stateDict missing key: ${key}`);
+                    return;
+                }
+                
+                const currentParam = p.param;
+                const is2D = Array.isArray(currentParam[0]);
+                
+                // Validate shape 
+                const savedRows = savedData.length;
+                const currentRows = currentParam.length;
+                
+                if (savedRows !== currentRows){
+                    console.warn(
+                        `[JST WARN]: stateDict shape mismatch for ${key} - ` +
+                        `saved rows: ${savedRows}, current rows: ${currentRows}. Skipping.`
+                    );
+                    return;
+                }
+                
+                if (is2D){
+                    const savedCols = savedData[0].length;
+                    const currentCols = currentParam[0].length;
+                    
+                    if (savedCols !== currentCols){
+                        console.warn(
+                            `[JST WARN]: stateDict shape mismatch for ${key} - ` +
+                            `saved cols: ${savedCols}, current cols: ${currentCols}. Skipping.`
+                        );
+                        return;
+                    }
+                    
+                    // Mutate in-place to preserve optimizer References
+                    for (let r=0; r<currentRows; r++){
+                        for (let c=0; c<currentCols; c++){
+                            currentParam[r][c] = savedData[r][c];
+                        }
+                    }
+                } else {
+                    // 1D parameter 
+                    for (let r=0; r<currentRows; r++){
+                        currentParam[r] = savedData[r];
+                    }
+                }
+            });
+            
+            // Invalidate cached flat representations
+            if(typeof layer._updateCache === 'function'){
+                layer._updateCache();
+            }
+        });
+        return this;
 	}
 }
 
@@ -1225,7 +1283,10 @@ export class CrossEntropyLoss{
 		return this._impl.forward(logits, targets);
 	}
 	backward(){
-		return this._impl.backward();
+		if (!this._impl.probs){
+            throw new Error("CrossEntropyLoss: backward() called before forward()");
+        }
+        return this._impl.backward();
 	}
 }
 
@@ -1959,40 +2020,76 @@ export class BatchNorm2d {
     eval() { this.training = false; }
 }
 
-// ---------------------- Model Save/Load (BETA) ----------------------
+// ---------------------- Model Save/Load (OPTIMIZED) ----------------------
+/**
+ * Serialize model parameters to a compact flat JSON string.
+ * Flattens all 2D parameter matrices into 1D arrays with shape metadata.
+*/
 export function saveModel(model){
-    if(!(model instanceof Sequential)){
+    if (!(model instanceof Sequential)){
         throw new Error("saveModel supports only Sequential models");
     }
     
-    const state = {
-        version: "2.0.0",
-        layers: model.layers.map((layer, idx) => {
-            const params = layer.parameters ? layer.parameters() : [];
+    const layers = [];
+    
+    for (let i=0; i<model.layers.length; i++){
+        const layer = model.layers[i];
+        const params = layer.parameters ? layer.parameters() : [];
+        const serializedParams = [];
+        
+        for (let j=0; j<params.length; j++){
+            const param = params[j].param;
+            const rows = param.length;
+            const is2D = Array.isArray(param[0]);
+            const cols = is2D ? param[0].length : 1;
             
-            if (params.length === 0){
-                return { type: layer.constructor.name, params: [] };
+            let flat;
+            
+            if (!is2D){
+                // 1D parameter: native slice for bulk memory copy 
+                flat = param.slice();
+            } else {
+                // 2D parameter: flatten row by row 
+                const total = rows * cols;
+                flat = new Array(total);
+                let cursor = 0;
+                for (let r=0; r<rows; r++){
+                    const row = param[r];
+                    for (let c=0; c<cols; c++){
+                        flat[cursor++] = row[c];
+                    }
+                }
             }
             
-            return {
-                type: layer.constructor.name,
-                params: params.map(p => ({
-                    // Deep clone parameter data 
-                    data: p.param.map(row =>
-                        Array.isArray(row) ? [...row] : row
-                    ),
-                    // Preserve shape metadata for validation 
-                    shape: Array.isArray(p.param[0])
-                        ? [p.param.length, p.param[0].length]
-                        : [p.param.length]
-                }))
+            serializedParams.push({
+                s: [rows, cols],
+                d: flat
+            });
+        }
+        
+        // Extract running stats for BatchNorm2d layers 
+        let runningStats = null;
+        if (typeof layer.runningMean !== 'undefined' && typeof layer.runningVar !== 'undefined'){
+            runningStats = {
+                mean: layer.runningMean.slice(),
+                var: layer.runningVar.slice()
             };
-        })
-    };
+        }
+        
+        layers.push({
+            t: layer.constructor.name,
+            p: serializedParams,
+            rs: runningStats
+        });
+    }
     
-    return JSON.stringify(state);
+    return JSON.stringify({ ver: "2.0.4", layers: layers });
 }
 
+/** 
+ * Load serialize flat parameters into a model.
+ * Uses pre-allocated arrays and single-cursor deserialization.
+*/
 export function loadModel(model, json){
     if (!(model instanceof Sequential)){
         throw new Error("loadModel supports only Sequential models");
@@ -2000,92 +2097,108 @@ export function loadModel(model, json){
     
     const state = JSON.parse(json);
     
-    // Validate structure 
     if (!state.layers || !Array.isArray(state.layers)){
         throw new Error("loadModel: invalid save format - missing 'layers' array");
     }
     
+    const layerCount = Math.min(state.layers.length, model.layers.length);
+    
     if (state.layers.length !== model.layers.length){
         console.warn(
             `[JST WARN]: Layer count mismatch - saved ${state.layers.length},` +
-            `current model has ${model.layers.length}. Loading what matches.`
+            `current model has ${model.layers.length}. Loading ${layerCount} layers.`
         );
     }
     
     let loadedCount = 0;
     let skippedCount = 0;
     
-    for(let i=0; i<Math.min(state.layers.length, model.layers.length); i++){
+    for (let i=0; i<layerCount; i++){
         const savedLayer = state.layers[i];
         const currentLayer = model.layers[i];
         
-        if(savedLayer.params.length === 0){
-            // Layer with no trainable params - skip 
-            continue 
+        if (!savedLayer.p || savedLayer.p.length === 0){
+            // Still restore running stats even if no trainable params 
+            if (savedLayer.rs && typeof currentLayer.runningMean !== 'undefined'){
+                for (let k=0; k<savedLayer.rs.mean.length; k++){
+                    currentLayer.runningMean[k] = savedLayer.rs.mean[k];
+                    currentLayer.runningVar[k] = savedLayer.rs.var[k];
+                }
+            }
+            continue;
         }
         
         // Validate layer type 
-        if (savedLayer.type !== currentLayer.constructor.name){
+        if (savedLayer.t !== currentLayer.constructor.name){
             console.warn(
                 `[JST WARN]: Layer ${i} type mismatch - ` +
-                `saved: ${savedLayer.type}, current: ${currentLayer.constructor.name}. Skipping.`
+                `saved: ${savedLayer.t}, current: ${currentLayer.constructor.name}. Skipping.`
             );
             skippedCount++;
             continue;
         }
         
-        // Get current layer parameters 
         const currentParams = currentLayer.parameters ? currentLayer.parameters() : [];
         
-        if (currentParams.length !== savedLayer.params.length){
+        if (currentParams.length !== savedLayer.p.length){
             console.warn(
                 `[JST WARN]: Layer ${i} parameter count mismatch - ` +
-                `saved: ${savedLayer.params.length}, current: ${currentParams.length}. Skipping.`
+                `saved: ${savedLayer.p.length}, current: ${currentParams.length}. Skipping.`
             );
             skippedCount++;
             continue;
         }
         
-        // Load parameters wiht shape validation 
-        for (let j=0; j<savedLayer.params.length; j++){
-            const savedParam = savedLayer.params[j];
-            const currentParam = currentParams[j].param;
+        for (let j=0; j<savedLayer.p.length; j++){
+            const savedParam = savedLayer.p[j];
+            const savedRows = savedParam.s[0];
+            const savedCols = savedParam.s[1];
+            const flatData = savedParam.d;
             
-            // Validate shape 
+            const currentParam = currentParams[j].param;
             const currentRows = currentParam.length;
             const currentCols = Array.isArray(currentParam[0])
                 ? currentParam[0].length 
                 : 1;
             
-            const savedRows = savedParam.shape[0];
-            const savedCols = savedParam.shape[1] || 1;
-            
+            // Validate shape 
             if (currentRows !== savedRows || currentCols !== savedCols){
                 console.warn(
                     `[JST WARN]: Layer ${i} param ${j} shape mismatch - ` +
-                    `saved: [${savedRows}, ${savedCols}],` +
-                    `current: [${currentRows}, ${currentCols}]. Skipping this parameter.`
+                    `saved: [${savedRows}, ${savedCols}], ` +
+                    `current: [${currentRows}, ${currentCols}]. Skipping.`
                 );
-                continue
+                continue;
             }
             
-            // Copy parameter data 
-            if (Array.isArray(currentParam[0])){
-                // 2D Parameter 
-                for (let r=0; r<currentRows; r++){
-                    for (let c=0; c<currentCols; c++){
-                        currentParam[r][c] = savedParam.data[r][c];
-                    }
+            // Single-cursor deserialization into pre-existing arrays 
+            let cursor =0;
+            
+            if (savedCols === 1 && !Array.isArray(currentParam[0])){
+                // 1D Parameter 
+                for (let r=0; r<savedRows; r++){
+                    currentParam[r] = flatData[cursor++];
                 }
             } else {
-                // 1D parameter 
-                for (let r=0; r<currentRows; r++){
-                    currentParam[r] = savedParam.data[r];
+                // 2D parameter: populate rows in-place
+                for (let r=0; r<savedRows; r++){
+                    const row = currentParam[r];
+                    for (let c=0; c<savedCols; c++){
+                        row[c] = flatData[cursor++];
+                    }
                 }
             }
         }
         
-        // Invalidate any cached flat representations 
+        // Restore BatchNorm2d running statistics 
+        if (savedLayer.rs && typeof currentLayer.runningMean !== 'undefined'){
+            for (let k=0; k<savedLayer.rs.mean.length; k++){
+                currentLayer.runningMean[k] = savedLayer.rs.mean[k];
+                currentLayer.runningVar[k] = savedLayer.rs.var[k];
+            }
+        }
+        
+        // Invalidate cached flat representations 
         if (typeof currentLayer._updateCache === 'function'){
             currentLayer._updateCache();
         }
@@ -2094,7 +2207,7 @@ export function loadModel(model, json){
     }
     
     console.log(
-        `[JST]: Model loaded: ${loadedCount} layers restored, ${skippedCount} skipped.`
+        `[JST] Model loaded: ${loadedCount} layers restored, ${skippedCount} skipped.`
     );
     
     return model;