mini-jstorch
Version:
A lightweight JavaScript neural network library for rapid frontend AI experimentation on low-resource devices Inspired by PyTorch.
255 lines (227 loc) • 13.7 kB
JavaScript
// MINI JSTORCH ENGINE - MAJOR ULTRA OPTIMIZED 1.2.3
// LICENSE: MIT (C) Rizal 2025
// ---------------------- Utilities ----------------------
function zeros(rows, cols) { return Array.from({length:rows},()=>Array(cols).fill(0)); }
function ones(rows, cols) { return Array.from({length:rows},()=>Array(cols).fill(1)); }
function randomMatrix(rows, cols, scale=0.1){ return Array.from({length:rows},()=>Array.from({length:cols},()=> (Math.random()*2-1)*scale)); }
function transpose(matrix){ return matrix[0].map((_,i)=>matrix.map(row=>row[i])); }
function addMatrices(a,b){ return a.map((row,i)=>row.map((v,j)=>v+(b[i] && b[i][j]!==undefined?b[i][j]:0))); }
function dot(a,b){ const res=zeros(a.length,b[0].length); for(let i=0;i<a.length;i++) for(let j=0;j<b[0].length;j++) for(let k=0;k<a[0].length;k++) res[i][j]+=a[i][k]*b[k][j]; return res; }
function softmax(x){ const m=Math.max(...x); const exps=x.map(v=>Math.exp(v-m)); const s=exps.reduce((a,b)=>a+b,0); return exps.map(v=>v/s); }
function crossEntropy(pred,target){ const eps=1e-12; return -target.reduce((sum,t,i)=>sum+t*Math.log(pred[i]+eps),0); }
// ---------------------- Tensor ----------------------
export class Tensor {
constructor(data){ this.data=data; this.grad=zeros(data.length,data[0].length); }
shape(){ return [this.data.length,this.data[0].length]; }
add(t){ return t instanceof Tensor?this.data.map((r,i)=>r.map((v,j)=>v+t.data[i][j])):this.data.map(r=>r.map(v=>v+t)); }
sub(t){ return t instanceof Tensor?this.data.map((r,i)=>r.map((v,j)=>v-t.data[i][j])):this.data.map(r=>r.map(v=>v-t)); }
mul(t){ return t instanceof Tensor?this.data.map((r,i)=>r.map((v,j)=>v*t.data[i][j])):this.data.map(r=>r.map(v=>v*t)); }
matmul(t){ if(t instanceof Tensor) return dot(this.data,t.data); else throw new Error("matmul requires Tensor"); }
transpose(){ return transpose(this.data); }
flatten(){ return this.data.flat(); }
static zeros(r,c){ return new Tensor(zeros(r,c)); }
static ones(r,c){ return new Tensor(ones(r,c)); }
static random(r,c,scale=0.1){ return new Tensor(randomMatrix(r,c,scale)); }
}
// ---------------------- Layers ----------------------
export class Linear {
constructor(inputDim,outputDim){
this.W=randomMatrix(inputDim,outputDim);
this.b=Array(outputDim).fill(0);
this.gradW=zeros(inputDim,outputDim);
this.gradb=Array(outputDim).fill(0);
this.x=null;
}
forward(x){
this.x=x;
const out=dot(x,this.W);
return out.map((row,i)=>row.map((v,j)=>v+this.b[j]));
}
backward(grad){
for(let i=0;i<this.W.length;i++) for(let j=0;j<this.W[0].length;j++)
this.gradW[i][j]=this.x.reduce((sum,row,k)=>sum+row[i]*grad[k][j],0);
for(let j=0;j<this.b.length;j++)
this.gradb[j]=grad.reduce((sum,row)=>sum+row[j],0);
const gradInput=zeros(this.x.length,this.W.length);
for(let i=0;i<this.x.length;i++)
for(let j=0;j<this.W.length;j++)
for(let k=0;k<this.W[0].length;k++)
gradInput[i][j]+=grad[i][k]*this.W[j][k];
return gradInput;
}
parameters(){ return [ {param:this.W,grad:this.gradW}, {param:[this.b],grad:[this.gradb]} ]; }
}
// ---------------------- Conv2D ----------------------
export class Conv2D {
constructor(inC,outC,kernel,stride=1,padding=0){
this.inC=inC; this.outC=outC; this.kernel=kernel;
this.stride=stride; this.padding=padding;
this.W=Array(outC).fill(0).map(()=>Array(inC).fill(0).map(()=>randomMatrix(kernel,kernel)));
this.gradW=Array(outC).fill(0).map(()=>Array(inC).fill(0).map(()=>zeros(kernel,kernel)));
this.x=null;
}
pad2D(input,pad){
return input.map(channel=>{
const rows=channel.length+2*pad;
const cols=channel[0].length+2*pad;
const out=Array.from({length:rows},()=>Array(cols).fill(0));
for(let i=0;i<channel.length;i++) for(let j=0;j<channel[0].length;j++) out[i+pad][j+pad]=channel[i][j];
return out;
});
}
conv2DSingle(input,kernel){
const rows=input.length-kernel.length+1;
const cols=input[0].length-kernel[0].length+1;
const out=zeros(rows,cols);
for(let i=0;i<rows;i++) for(let j=0;j<cols;j++)
for(let ki=0;ki<kernel.length;ki++) for(let kj=0;kj<kernel[0].length;kj++)
out[i][j]+=input[i+ki][j+kj]*kernel[ki][kj];
return out;
}
forward(batch){
this.x=batch;
return batch.map(sample=>{
const channelsOut=[];
for(let oc=0;oc<this.outC;oc++){
let outChan=zeros(sample[0].length,sample[0][0].length);
for(let ic=0;ic<this.inC;ic++){
let inputChan=sample[ic];
if(this.padding>0) inputChan=this.pad2D([inputChan],this.padding)[0];
const conv=this.conv2DSingle(inputChan,this.W[oc][ic]);
outChan=addMatrices(outChan,conv);
}
channelsOut.push(outChan);
}
return channelsOut;
});
}
backward(grad) {
const batchSize = this.x.length;
const gradInput = this.x.map(sample => sample.map(chan => zeros(chan.length, chan[0].length)));
const gradW = this.W.map(oc => oc.map(ic => zeros(this.kernel,this.kernel)));
for (let b = 0; b < batchSize; b++) {
const xPadded = this.pad2D(this.x[b], this.padding);
const gradInputPadded = xPadded.map(chan => zeros(chan.length, chan[0].length));
for (let oc = 0; oc < this.outC; oc++) {
for (let ic = 0; ic < this.inC; ic++) {
const outGrad = grad[b][oc];
const inChan = xPadded[ic];
// Compute gradW
for (let i = 0; i < this.kernel; i++) {
for (let j = 0; j < this.kernel; j++) {
let sum = 0;
for (let y = 0; y < outGrad.length; y++) {
for (let x = 0; x < outGrad[0].length; x++) {
const inY = y * this.stride + i;
const inX = x * this.stride + j;
if (inY < inChan.length && inX < inChan[0].length) {
sum += inChan[inY][inX] * outGrad[y][x];
}
}
}
gradW[oc][ic][i][j] += sum;
}
}
// Compute gradInput
const flippedKernel = this.W[oc][ic].map(row => [...row].reverse()).reverse();
for (let y = 0; y < outGrad.length; y++) {
for (let x = 0; x < outGrad[0].length; x++) {
for (let i = 0; i < this.kernel; i++) {
for (let j = 0; j < this.kernel; j++) {
const inY = y * this.stride + i;
const inX = x * this.stride + j;
if (inY < gradInputPadded[ic].length && inX < gradInputPadded[ic][0].length) {
gradInputPadded[ic][inY][inX] += flippedKernel[i][j] * outGrad[y][x];
}
}
}
}
}
}
}
// Remove padding from gradInput
if (this.padding > 0) {
for (let ic = 0; ic < this.inC; ic++) {
const padded = gradInputPadded[ic];
const cropped = padded.slice(this.padding, padded.length - this.padding)
.map(row => row.slice(this.padding, row.length - this.padding));
gradInput[b][ic] = cropped;
}
} else {
for (let ic = 0; ic < this.inC; ic++) gradInput[b][ic] = gradInputPadded[ic];
}
}
this.gradW = gradW;
return gradInput;
}
parameters(){ return this.W.flatMap((w,oc)=>w.map((wc,ic)=>({param:wc,grad:this.gradW[oc][ic]}))); }
}
// ---------------------- Sequential ----------------------
export class Sequential {
constructor(layers=[]){ this.layers=layers; }
forward(x){ return this.layers.reduce((acc,l)=>l.forward(acc), x); }
backward(grad){ return this.layers.reduceRight((g,l)=>l.backward(g), grad); }
parameters(){ return this.layers.flatMap(l=>l.parameters?l.parameters():[]); }
}
// ---------------------- Activations ----------------------
export class ReLU{ constructor(){ this.out=null; } forward(x){ this.out=x.map(r=>r.map(v=>Math.max(0,v))); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*(this.out[i][j]>0?1:0))); } }
export class Sigmoid{ constructor(){ this.out=null; } forward(x){ const fn=v=>1/(1+Math.exp(-v)); this.out=x.map(r=>r.map(fn)); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*this.out[i][j]*(1-this.out[i][j]))); } }
export class Tanh{ constructor(){ this.out=null; } forward(x){ this.out=x.map(r=>r.map(v=>Math.tanh(v))); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*(1-this.out[i][j]**2))); } }
export class LeakyReLU{ constructor(alpha=0.01){ this.alpha=alpha; this.out=null; } forward(x){ this.out=x.map(r=>r.map(v=>v>0?v:v*this.alpha)); return this.out; } backward(grad){ return grad.map((r,i)=>r.map((v,j)=>v*(this.out[i][j]>0?1:this.alpha))); } }
export class GELU{ constructor(){ this.out=null; } forward(x){ const fn=v=>0.5*v*(1+Math.tanh(Math.sqrt(2/Math.PI)*(v+0.044715*v**3))); this.out=x.map(r=>r.map(fn)); return this.out; } backward(grad){ return grad.map((r,i)=>r.map(v=>v*1)); } }
// ---------------------- Dropout ----------------------
export class Dropout{ constructor(p=0.5){ this.p=p; } forward(x){ return x.map(r=>r.map(v=>v*Math.random()>=this.p?v:0)); } backward(grad){ return grad.map(r=>r.map(v=>v*(1-this.p))); } }
// ---------------------- Losses ----------------------
export class MSELoss{ forward(pred,target){ this.pred=pred; this.target=target; const losses=pred.map((row,i)=>row.reduce((sum,v,j)=>sum+(v-target[i][j])**2,0)/row.length); return losses.reduce((a,b)=>a+b,0)/pred.length; } backward(){ return this.pred.map((row,i)=>row.map((v,j)=>2*(v-this.target[i][j])/row.length)); } }
export class CrossEntropyLoss{ forward(pred,target){ this.pred=pred; this.target=target; const losses=pred.map((p,i)=>crossEntropy(softmax(p),target[i])); return losses.reduce((a,b)=>a+b,0)/pred.length; } backward(){ return this.pred.map((p,i)=>{ const s=softmax(p); return s.map((v,j)=>(v-this.target[i][j])/this.pred.length); }); } }
// ---------------------- Optimizers ----------------------
export class Adam{
constructor(params,lr=0.001,b1=0.9,b2=0.999,eps=1e-8){
this.params=params; this.lr=lr; this.beta1=b1; this.beta2=b2; this.eps=eps;
this.m=params.map(p=>zeros(p.param.length,p.param[0].length||1));
this.v=params.map(p=>zeros(p.param.length,p.param[0].length||1));
this.t=0;
}
step(){
this.t++;
this.params.forEach((p,idx)=>{
for(let i=0;i<p.param.length;i++)
for(let j=0;j<(p.param[0].length||1);j++){
const g=p.grad[i][j];
this.m[idx][i][j]=this.beta1*this.m[idx][i][j]+(1-this.beta1)*g;
this.v[idx][i][j]=this.beta2*this.v[idx][i][j]+(1-this.beta2)*g*g;
const mHat=this.m[idx][i][j]/(1-Math.pow(this.beta1,this.t));
const vHat=this.v[idx][i][j]/(1-Math.pow(this.beta2,this.t));
p.param[i][j]-=this.lr*mHat/(Math.sqrt(vHat)+this.eps);
}
});
}
}
export class SGD{ constructor(params,lr=0.01){ this.params=params; this.lr=lr; } step(){ this.params.forEach(p=>{ for(let i=0;i<p.param.length;i++) for(let j=0;j<(p.param[0].length||1);j++) p.param[i][j]-=this.lr*p.grad[i][j]; }); } }
// ---------------------- Model Save/Load ----------------------
export function saveModel(model){
if(!(model instanceof Sequential)) throw new Error("saveModel supports only Sequential");
const weights=model.layers.map(layer=>({weights:layer.W||null,biases:layer.b||null}));
return JSON.stringify(weights);
}
export function loadModel(model,json){
if(!(model instanceof Sequential)) throw new Error("loadModel supports only Sequential");
const weights=JSON.parse(json);
model.layers.forEach((layer,i)=>{
if(layer.W && weights[i].weights) layer.W=weights[i].weights;
if(layer.b && weights[i].biases) layer.b=weights[i].biases;
});
}
// ---------------------- Advanced Utils ----------------------
export function flattenBatch(batch){ return batch.flat(2); }
export function stack(tensors){ return tensors.map(t=>t.data); }
export function eye(n){ return Array.from({length:n},(_,i)=>Array.from({length:n},(_,j)=>i===j?1:0)); }
export function concat(a,b,axis=0){ /* concat along axis */ if(axis===0) return [...a,...b]; if(axis===1) return a.map((row,i)=>[...row,...b[i]]); }
export function reshape(tensor, rows, cols) {
let flat = tensor.data.flat(); // flatten dulu
if(flat.length < rows*cols) throw new Error("reshape size mismatch");
const out = Array.from({length: rows}, (_, i) =>
flat.slice(i*cols, i*cols + cols)
);
return out;
}