Generic Compute Neural Network (Modular)

/*
    Neural networks - tested with xor and back propagation

    Array 'Blocks' for the Neural Network Data (layer outputs, weights, errors, ..)

    CPU and GPU VERSION - Compares/Checks Results

    Modular funtions - load/save/reset/ (both cpu/gpu) - compare
*/

console.log('xnet.js');

let xlayers       = [2,3,1];
let xbuild        = 'cpu';
let xiterations   = 1;
let xlearningrate = 0.1;

xinitialize = ( arg = { layers:[1,2,1], build:'cpu', iterations:10, learningrate:0.2 } ) => {};
xactivate   = async ( ip )  => { };
xpropagate  = async ( op )  => { };

//--------------------------------------------------------------------------------------

const config = {
    layers: [2, 3, 2], // Number of neurons per layer (input layer, hidden layers, output layer)

    weights: [
      // Weights between layer 0 and layer 1
      [
        [0.5, -0.2, 0.1], // Weights from neuron 0 in layer 0 to neurons in layer 1
        [0.3, 0.8, -0.5],
      ],
      // Weights between layer 1 and layer 2
      [
        [0.1, -0.4], // Weights from neuron 0 in layer 1 to neurons in layer 2
        [-0.7, 0.8],
        [0.6,  0.2],
      ]
    ],
    biases: [
      // Biases for layer 0
      [0.0, 0.0],
      // Biases for layer 1
      [0.1, -0.2, 0.3],
      // Biases for layer 2
      [0.5, -0.1]
    ],
    // test data to check the network is connected correctly with coefficients
    check: {
      input:  [ 0, 1 ],
      output: [ 0.5933927893638611, 0.5663766860961914 ]
    }
};

//--------------------------------------------------------------------------------------

let   LEARNING_RATE = 0.2;
const VARIANCE_W    = 0.5;
const randomUniform = (min, max) => Math.random() * (max - min) + min;

const ru = ()=>{ return randomUniform(-VARIANCE_W, VARIANCE_W); }

let   layers  = [ 2, 3, 1 ];
let   maxLayerSize = [...layers].sort( (a,b)=>b-a )[0];

const xordataset = [  { inputs: [0,0], outputs: [0] },
                      { inputs: [0,1], outputs: [1] },
                      { inputs: [1,0], outputs: [1] },
                      { inputs: [1,1], outputs: [0] }  ];

console.assert( xordataset[0].outputs.length == layers[ layers.length-1 ] );

let   weights  = Array( layers.length * maxLayerSize * maxLayerSize ).fill( 0 );
let   biases   = Array( layers.length * maxLayerSize ).fill(0);
let   loutputs = Array( layers.length * maxLayerSize ).fill(0);
let   errors   = Array( layers.length * maxLayerSize ).fill(0);

let   MAX_NEURONS_PER_LAYER = maxLayerSize;
let   NUM_LAYERS            = layers.length;

//--------------------------------------------------------------------------------------

// Default is random weights and biases - overridden later on by config file if loaded
const initializeWeightsandBiases = () => {
      console.log('initializeWeightsandBiases');
    for (let i=0; i<layers.length-1; i++)
    {
       for (let k=0; k<layers[i]; i++)
       {
            for (let g=0; g<layers[i+1]; g++)
            {
                setWeight( i, k, g , ru() );
            }
       }
    }

    for (let i=0; i<layers.length-1; i++)
    {
       for (let k=0; k<layers[i]; i++)
       {
             setBias( i, k, 0.0 );
       }
    }
}
initializeWeightsandBiases();

//--------------------------------------------------------------------------------------

function getBias(layer, neuron) {
  return biases[layer * MAX_NEURONS_PER_LAYER + neuron];
}

function setBias(layer, neuron, value) {
  biases[layer * MAX_NEURONS_PER_LAYER + neuron] = value;
}

function setOutput(layer, neuron, value) {
  loutputs[layer * MAX_NEURONS_PER_LAYER + neuron] = value;
}

function getOutput(layer, neuron) {
  return loutputs[layer * MAX_NEURONS_PER_LAYER + neuron];
}

function getWeight(layer, fromNeuron, toNeuron) {
  return weights[layer * MAX_NEURONS_PER_LAYER * MAX_NEURONS_PER_LAYER
                + fromNeuron * MAX_NEURONS_PER_LAYER
                + toNeuron];
}

function setWeight(layer, fromNeuron, toNeuron, value) {
  weights[layer * MAX_NEURONS_PER_LAYER * MAX_NEURONS_PER_LAYER
         + fromNeuron * MAX_NEURONS_PER_LAYER
         + toNeuron] = value;
}

function setError(layer, neuron, value) {
  errors[layer * MAX_NEURONS_PER_LAYER + neuron] = value;
}

function getError(layer, neuron) {
  return errors[layer * MAX_NEURONS_PER_LAYER + neuron];
}

//--------------------------------------------------------------------------------------

function parseConfig( json )
{
    //if ( typeof(config.weights) != 'undefined' )
    for (let i=0; i<config.weights.length; i++)
    {
       const layer0 = i;
       const layer1 = i+1;

       const layer0N = config.weights[i].length;
       const layer1N = config.weights[i][0].length;

       console.assert( layer0N == config.layers[i] );
       console.assert( layer1N == config.layers[i+1] );

       for (let n0=0; n0<layer0N; n0++)
       {
          for (let n1=0; n1<layer1N; n1++)
          {
              const weight  = config.weights[layer0][n0][n1];
              setWeight( layer0, n0, n1, weight );
          }
       }
    }

    // setup biases
    //if ( typeof(config.biases) != 'undefined' )
    for (let b=0; b<config.biases.length; b++)
    {
       // layer
       for (let i=0; i<config.biases[b].length; i++)
       {
           // each neuron
           const bias = config.biases[b][i];
           setBias( b, i, bias );
       }
    }
}

//--------------------------------------------------------------------------------------

function loadConfig()
{
    let inputdiv    = document.createElement('input');
    inputdiv.type   = 'file';
    inputdiv.accept = 'json/json';
    inputdiv.onchange = function()
    {
          const fileList = this.files;
          let reader = new FileReader(); // no arguments

          reader.readAsText( this.files[0] );

          reader.onload = function() {

        parseConfig( reader.result );
      };
    }
    inputdiv.click();
}

//--------------------------------------------------------------------------------------

function saveConfig()
{
    // Save config network (layout, weights, biases, ..)

    let newconfig = {
      layers  : [],
      weights : [],
      biases  : [],
      check   : {
              input : [],
              output: [] }
    };

    // layers
    for (let i=0; i<layers.length; i++)
    {
       newconfig.layers.push( layers[i] );
    }

    // weights
    for (let i=0; i<layers.length-1; i++)
    {
        let layerweights = [];
        for (let n0=0; n0<layers[i]; n0++)
        {
          let weights = [];
          for (let n1=0; n1<layers[i+1]; n1++)
          {
                 const weight = getWeight( i, n0, n1 );
              weights.push( weight );
          }
          layerweights.push( weights );
        }
        newconfig.weights.push( layerweights );
    }

    // biases
    for (let i=0; i<layers.length; i++)
    {
        let layerbiases = [];
        for (let n0=0; n0<layers[i]; n0++)
        {
            const bias = getBias( i, n0 );
            layerbiases.push( bias );
        }
        newconfig.biases.push( layerbiases );
    }

    //newconfig.check.input  = dataset[1].inputs;
    //newconfig.check.output = activate( dataset[1].inputs );

    // download dialog
    let json = JSON.stringify( newconfig );
    let blob2 = new Blob([json], {type: "text/txt"});
    let url2 = window.URL.createObjectURL(blob2);
    let link = document.createElement('a');
    document.body.appendChild( link );
    link.innerHTML = 'download config';
    link.download= 'nnconfig.json';
    link.href = url2 ;
    link.click();
}

//--------------------------------------------------------------------------------------

const sigmoid = (x) => 1.0 / (1.0 + Math.exp(-x));

const sigmoidDerivative = (x) => x * (1 - x);

const relu = (x) =>  { return Math.max(0.0, x); }

const reluDerivative = (x) => {  if (x > 0.0) { return 1.0; } return 0.0; }

const leakyRelu = (x, alpha = 0.01) => { return x > 0 ? x : alpha * x; }

const leakyReluDerivative = (x, alpha = 0.01) => { return x > 0 ? 1 : alpha; }

const activate = (iin) => {
      //console.assert( iin.length == outputs[ 0 ].length );
      //console.assert( weights[0].length    == layers[0] );
      //console.assert( weights[0][0].length == layers[1] );
      //console.assert( weights[1].length    == layers[1] );
    //console.assert( weights[1][0].length == layers[2] );

    for (let i=0; i<NUM_LAYERS; i++)
    {
          if ( i==0 )
        {
              for (let k=0; k<iin.length; k++)
            {
                setOutput(0, k, iin[ k ] );
            }
        }
        else
        {
              for (let k=0; k<layers[i]; k++)
            {
                  var sum = 0.0;
                  for (let b=0; b<layers[i-1]; b++)
                {
                      sum += getOutput( i-1, b ) * getWeight( i-1, b, k );
                }
                setOutput( i, k,  sigmoid( sum + getBias( i, k ) ) );
            }

        }
    }

      let out = [];
    for (let k=0; k<layers[ NUM_LAYERS-1 ]; k++)
    {
        out.push( getOutput( NUM_LAYERS-1, k ) );
    }
    return out;
};

//--------------------------------------------------------------------------------------

const propagate = ( target ) => {

    for (let i=NUM_LAYERS-1; i>0; i--)
    {
          for (let k=0; k<layers[i]; k++)
        {
              if ( i==NUM_LAYERS-1 )
            {
                  let error = ( target[k] - getOutput( i, k ) ) * sigmoidDerivative( getOutput( i,k ) );
                setError( i, k, error );
            }
            else
            {
                  setError( i, k, 0.0 );
                  for (let g=0; g<layers[i+1]; g++)
                {
                      let error = getError( i+1, g ) * getWeight( i,k,g ) * sigmoidDerivative( getOutput(i,k) );
                    setError( i, k, error );
                }
            }
        }
    }

    for (let i=0; i<NUM_LAYERS; i++)
    {
          for (let k=0; k<layers[i]; k++)
        {
            if ( i < NUM_LAYERS-1 ) {
              for (let g=0; g<layers[i+1]; g++)
            {
                  var weight = getWeight( i, k, g );
                  weight += LEARNING_RATE * getOutput(i,k) * getError(i+1,g);
                  setWeight( i, k, g, weight );
            }
            }

            let bias = getBias( i,k );
            bias += LEARNING_RATE * getError( i, k );
            setBias( i, k, bias );
        }
    }
};

//--------------------------------------------------------------------------------------

console.log( new Date() );

//--------------------------------------------------------------------------------------

function outputJSNetwork()
{
    let cost = 0;
    for (let j = 0; j < xordataset.length; j++) {
      let o = activate(  xordataset[j].inputs );
      for (let b=0; b<xordataset[j].outputs.length; b++)
      {
        cost += Math.pow( xordataset[j].outputs[b] - o[b], 2);
      }
    }
    cost /= 4;
    console.log(`total error (no training) mean squared error: ${cost}`);

    for (let i = 0; i < xordataset.length; i++)
    {
        const result = activate( xordataset[i].inputs );

        console.log(`for input ${xordataset[i].inputs} expected ${xordataset[i].outputs} predicted ${result[0].toFixed(4)} which is ${Math.round(result[0]) === xordataset[i].outputs[0] ? "correct" : "incorrect"}`);
    }

}

//--------------------------------------------------------------------------------------

function trainJSNetwork()
{
    console.log(' ----------TRAIN WITH GPU VERSION------------- ' );
    // - test the neural network using iteration loop - xor dataset

    for (let epoch = 0; epoch <= 10000; epoch++) {
        let indexes = Array.from(Array( xordataset.length ).keys());
        indexes.sort(() => Math.random() - 0.5);
        for (let j of indexes) {
            activate( xordataset[j].inputs );
            propagate( xordataset[j].outputs, LEARNING_RATE);
        }

        if (epoch % 1000 === 0) {
            let cost = 0;
            for (let j = 0; j < xordataset.length; j++) {
                let o = activate(  xordataset[j].inputs );
                for (let b=0; b<xordataset[j].outputs.length; b++)
                {
                    cost += Math.pow( xordataset[j].outputs[b] - o[b], 2);
                }
            }
            cost /= 4;
            console.log(`epoch ${epoch} mean squared error: ${cost}`);
        }
    }

    for (let i = 0; i < xordataset.length; i++)
    {
        const result = activate( xordataset[i].inputs );

        console.log(`for input ${xordataset[i].inputs} expected ${xordataset[i].outputs} predicted ${result[0].toFixed(4)} which is ${Math.round(result[0]) === xordataset[i].outputs[0] ? "correct" : "incorrect"}`);
    }

}

//--------------------------------------------------------------------------------------

const but = ( s, n ) =>
{
  let butEl = document.createElement('button');
  butEl.onclick = () => { n(); }
  butEl.innerHTML = s;
  document.body.appendChild( butEl );
  document.body.appendChild( document.createElement('br') );
}

//---------------------------------------------------------------------------

const adapter = await navigator.gpu.requestAdapter();
const device  = await adapter.requestDevice();

//---------------------------------------------------------------------------

async function runComputePipeline(shaderCode, bindings, workGroupCount, entryFunction='main') {
    const shaderModule = device.createShaderModule({ code: shaderCode });

    let entries = [];
    for (let n=0; n<bindings.length; n++) {
       entries.push( {binding: bindings[n].binding, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"} } );
    }
    const bindGroupLayout = device.createBindGroupLayout( { "entries": entries } );
      /*
    const bindGroupLayout = device.createBindGroupLayout({
        entries: [ {binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  },
                   {binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  },
                   {binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  },
                   {binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"}  }
                 ]
      });
    */

    const pipeline = device.createComputePipeline({
      layout: device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]}),
      compute: { module: shaderModule, entryPoint: entryFunction },
    });

    const bindGroup = device.createBindGroup({ layout: bindGroupLayout, entries: bindings });

    const commandEncoder = device.createCommandEncoder();
    const passEncoder = commandEncoder.beginComputePass();
    passEncoder.setPipeline(pipeline);
    passEncoder.setBindGroup(0, bindGroup);
    passEncoder.dispatchWorkgroups(workGroupCount);
    passEncoder.end();

    const commandBuffer = commandEncoder.finish();
    device.queue.submit([commandBuffer]);
    await device.queue.onSubmittedWorkDone();

    // Return pipeline/group
    return { pipeline:pipeline, bindGroup:bindGroup, workGroupCount:workGroupCount };
}

function createBuffer( data, usage ) {
    const buffer = device.createBuffer({
      size: data.byteLength,
      usage: usage | GPUBufferUsage.COPY_DST,
      mappedAtCreation: true,
    });
    new data.constructor(buffer.getMappedRange()).set(data);
    buffer.unmap();
    return buffer;
}

let stConfig = {
    MAX_NEURONS_PER_LAYER : maxLayerSize,  // u32,
    NUM_LAYERS            : NUM_LAYERS,    // u32,
    LEARNING_RATE         : LEARNING_RATE*1000000, // f32, // 0.1, 0.2, ...
    ACTIVATION_TYPE       : 0              // u32  // enum - 0-sig, 1-relu, 2-reluleaky, ...
}

let configBuffer    = createBuffer(new Uint32Array( Object.values(stConfig) ), GPUBufferUsage.STORAGE );
let layersBuffer    = createBuffer(new Uint32Array( layers ), GPUBufferUsage.STORAGE );
let inputsBuffer    = createBuffer(new Float32Array( layers[0]  ), GPUBufferUsage.STORAGE );
let resultsBuffer   = createBuffer(new Float32Array( layers[layers.length-1] ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC );
let expectedBuffer  = createBuffer(new Float32Array( layers[layers.length-1] ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC );
let weightsBuffer   = createBuffer(new Float32Array( weights.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
let biasesBuffer    = createBuffer(new Float32Array( biases.flat(4)  ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
let outputsBuffer   = createBuffer(new Float32Array( loutputs.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
let errorsBuffer    = createBuffer(new Float32Array( errors.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);

//---------------------------------------------------------------------------

// Shaders
const forwardShaderCode = `
//const MAX_NEURONS_PER_LAYER = ${maxLayerSize};
//const NUM_LAYERS            = ${NUM_LAYERS};

struct stConfig {
    MAX_NEURONS_PER_LAYER : u32,
    NUM_LAYERS            : u32,
    LEARNING_RATE         : u32, // 0.1, 0.2, ...
    ACTIVATION_TYPE       : u32  // enum 0-sig, 1-relu, 2-reluleaky, ...
}

@group(0) @binding(0) var<storage, read_write> layers       : array<u32>;
@group(0) @binding(1) var<storage, read_write> weights      : array<f32>;
@group(0) @binding(2) var<storage, read_write> biases       : array<f32>;
@group(0) @binding(3) var<storage, read_write> inputs       : array<f32>;
@group(0) @binding(4) var<storage, read_write> results      : array<f32>;
@group(0) @binding(5) var<storage, read_write> outputs      : array<f32>;
@group(0) @binding(6) var<storage, read_write> config       : stConfig;

// Activation function (Sigmoid)
fn sigmoid(x:f32) -> f32 {
  return 1.0 / (1.0 + exp(-x));
}

fn getBias(layer: u32, neuron: u32) -> f32 {
    return biases[ layer * config.MAX_NEURONS_PER_LAYER + neuron ];
}

fn setOutput(layer: u32, neuron: u32, value: f32) {
    outputs[ layer * config.MAX_NEURONS_PER_LAYER + neuron ] = value;
}

fn getOutput(layer: u32, neuron: u32) -> f32 {
    return outputs[ layer * config.MAX_NEURONS_PER_LAYER + neuron ];
}

fn getWeight(layer: u32, fromNeuron: u32, toNeuron: u32 ) -> f32 {
    return weights[ layer * config.MAX_NEURONS_PER_LAYER * config.MAX_NEURONS_PER_LAYER
                   + fromNeuron * config.MAX_NEURONS_PER_LAYER
                   + toNeuron ];
}

@compute @workgroup_size( 1 )
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    // Single threaded version - basic compute testing reconisonse

    for (var i:u32=0; i<config.NUM_LAYERS; i++)
    {
          if ( i==0 )
        {
              for (var k:u32=0; k< arrayLength(&inputs); k++)
            {
                setOutput(0, k, inputs[ k ] );
            }
        }
        else
        {
              for (var k:u32=0; k<layers[i]; k++)
            {
                  var sum = 0.0;
                  for (var b:u32=0; b<layers[i-1]; b++)
                {
                      sum += getOutput( i-1, b ) * getWeight( i-1, b, k );
                }
                setOutput( i, k,  sigmoid( sum + getBias( i, k ) ) );
            }
        }
    }

    for (var k:u32=0; k< arrayLength(&results); k++)
    {
        results[k] = getOutput( arrayLength(&layers)-1 , k );
    }
}
`;

async function activateGPU( inputs )
{
    device.queue.writeBuffer( inputsBuffer,     0, new Float32Array( inputs ) );
    await device.queue.onSubmittedWorkDone();

    await runComputePipeline(forwardShaderCode, [
                  { binding: 0, resource: { buffer: layersBuffer        } },
                  { binding: 1, resource: { buffer: weightsBuffer       } },
                  { binding: 2, resource: { buffer: biasesBuffer        } },
                  { binding: 3, resource: { buffer: inputsBuffer        } },
                  { binding: 4, resource: { buffer: resultsBuffer       } },
                  { binding: 5, resource: { buffer: outputsBuffer       } },
                  { binding: 6, resource: { buffer: configBuffer        } },
            ], 1 );

    // Retrieve results
    const readBuffer = device.createBuffer({size: resultsBuffer.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });

    const commandEncoder = device.createCommandEncoder();
    commandEncoder.copyBufferToBuffer(resultsBuffer, 0, readBuffer, 0, resultsBuffer.size);
    device.queue.submit([commandEncoder.finish()]);

    await readBuffer.mapAsync(GPUMapMode.READ);
    const results = Array.from( new Float32Array(readBuffer.getMappedRange()) );
    readBuffer.unmap();
      return results;
}

//---------------------------------------------------------------------------

// Shaders
const backwardShaderCode = `
//const MAX_NEURONS_PER_LAYER = ${maxLayerSize};
//const NUM_LAYERS            = ${NUM_LAYERS};
//const LEARNING_RATE:f32     = ${LEARNING_RATE};

struct stConfig {
    MAX_NEURONS_PER_LAYER : u32,
    NUM_LAYERS            : u32,
    LEARNING_RATE         : u32, // 0.1, 0.2, ...
    ACTIVATION_TYPE       : u32  // enum 0-sig, 1-relu, 2-reluleaky, ...
}

@group(0) @binding(0) var<storage, read_write> layers       : array<u32>;
@group(0) @binding(1) var<storage, read_write> weights      : array<f32>;
@group(0) @binding(2) var<storage, read_write> biases       : array<f32>;
@group(0) @binding(3) var<storage, read_write> inputs       : array<f32>;
@group(0) @binding(4) var<storage, read_write> outputs      : array<f32>;
@group(0) @binding(5) var<storage, read_write> expected     : array<f32>;
@group(0) @binding(6) var<storage, read_write> errors       : array<f32>;
@group(0) @binding(7) var<storage, read_write> config       : stConfig;

// Derivative of sigmoid function
fn sigmoidDerivative(x:f32) -> f32 {
  return x * (1 - x);
}

fn getBias(layer: u32, neuron: u32) -> f32 {
    return biases[ layer * config.MAX_NEURONS_PER_LAYER + neuron ];
}

fn setBias(layer: u32, neuron: u32, value:f32 ) {
    biases[ layer * config.MAX_NEURONS_PER_LAYER + neuron ] = value;
}

fn setOutput(layer: u32, neuron: u32, value: f32) {
    outputs[ layer * config.MAX_NEURONS_PER_LAYER + neuron ] = value;
}

fn getOutput(layer: u32, neuron: u32) -> f32 {
    return outputs[ layer * config.MAX_NEURONS_PER_LAYER + neuron ];
}

fn getWeight(layer: u32, fromNeuron: u32, toNeuron: u32 ) -> f32 {
    return weights[ layer * config.MAX_NEURONS_PER_LAYER * config.MAX_NEURONS_PER_LAYER
                   + fromNeuron * config.MAX_NEURONS_PER_LAYER
                   + toNeuron ];
}

fn setWeight(layer: u32, fromNeuron: u32, toNeuron: u32, value:f32 ) {
    weights[ layer * config.MAX_NEURONS_PER_LAYER * config.MAX_NEURONS_PER_LAYER
                   + fromNeuron * config.MAX_NEURONS_PER_LAYER
                   + toNeuron ] = value;
}

fn setError( layer: u32, neuron: u32, value: f32 ){
    errors[ layer * config.MAX_NEURONS_PER_LAYER + neuron ] = value;
}

fn getError( layer: u32, neuron: u32 ) -> f32 {
    return errors[ layer * config.MAX_NEURONS_PER_LAYER + neuron ];
}

@compute @workgroup_size( 1 )
fn main(@builtin(global_invocation_id) global_id: vec3<u32>)
{
    // Single threaded - test compute reconisone

    let LEARNING_RATE:f32 = f32( config.LEARNING_RATE )/1000000;

    for (var i:u32=config.NUM_LAYERS-1; i>0; i--)
    {
          for (var k:u32=0; k<layers[i]; k++)
        {
              if ( i==config.NUM_LAYERS-1 )
            {
                  let error = ( expected[k] - getOutput( i, k ) ) * sigmoidDerivative( getOutput( i,k ) );
                setError( i, k, error );
            }
            else
            {
                  setError( i, k, 0.0 );
                  for (var g:u32=0; g<layers[i+1]; g++)
                {
                      let error = getError( i+1, g ) * getWeight( i,k,g ) * sigmoidDerivative( getOutput(i,k) );
                    setError( i, k, error );
                }
            }
        }
    }


    for (var i:u32=0; i<config.NUM_LAYERS; i++)
    {
          for (var k:u32=0; k<layers[i]; k++)
        {
            if ( i < config.NUM_LAYERS-1 ) {
              for (var g:u32=0; g<layers[i+1]; g++)
            {
                  var weight = getWeight( i, k, g );
                  weight += LEARNING_RATE * getOutput(i,k) * getError(i+1,g);
                  setWeight( i, k, g, weight );
            } }

            var bias = getBias( i,k );
            bias += LEARNING_RATE * getError( i, k );
            setBias( i, k, bias );
        }
    }
}
`;

async function propagateGPU( expected )
{
    device.queue.writeBuffer( expectedBuffer,   0, new Float32Array( expected ) );
    await device.queue.onSubmittedWorkDone();

    await runComputePipeline(backwardShaderCode, [
            { binding: 0, resource: { buffer: layersBuffer       } },
          { binding: 1, resource: { buffer: weightsBuffer      } },
          { binding: 2, resource: { buffer: biasesBuffer       } },
          { binding: 3, resource: { buffer: inputsBuffer       } },
          { binding: 4, resource: { buffer: outputsBuffer      } },
          { binding: 5, resource: { buffer: expectedBuffer     } },
          { binding: 6, resource: { buffer: errorsBuffer       } },
          { binding: 7, resource: { buffer: configBuffer        } },
    ], 1 );
}

//---------------------------------------------------------------------------

async function outputGPUNetwork()
{
    console.log(`GPU VERSION (should match CPU version) `);

    device.queue.writeBuffer( weightsBuffer,     0, new Float32Array( weights.flat(4) ) );
    device.queue.writeBuffer( biasesBuffer,      0, new Float32Array( biases.flat(4)  ) );
    await device.queue.onSubmittedWorkDone();

      let cost = 0;
    for (let j = 0; j < xordataset.length; j++) {
      let o = await activateGPU(  xordataset[j].inputs );
      for (let b=0; b<xordataset[j].outputs.length; b++)
      {
        cost += Math.pow( xordataset[j].outputs[b] - o[b], 2);
      }
    }
    cost /= 4;
    console.log(`total error (no training) mean squared error: ${cost}`);


    for (let i = 0; i < xordataset.length; i++)
    {
        const result = await activateGPU( xordataset[i].inputs );

        console.log(`for input ${xordataset[i].inputs} expected ${xordataset[i].outputs} predicted ${result[0].toFixed(4)} which is ${Math.round(result[0]) === xordataset[i].outputs[0] ? "correct" : "incorrect"}`);
    }
}

//---------------------------------------------------------------------------

async function trainGPUNetwork()
{
    console.log(`
    ----------TRAIN WITH GPU VERSION-------------
    `);

    device.queue.writeBuffer( weightsBuffer,     0, new Float32Array( weights.flat(4) ) );
    device.queue.writeBuffer( biasesBuffer,      0, new Float32Array( biases.flat(4)  ) );
    await device.queue.onSubmittedWorkDone();

    for (let epoch = 0; epoch <= 1000; epoch++) {
        let indexes = Array.from(Array( xordataset.length ).keys());
        indexes.sort(() => Math.random() - 0.5);
        for (let j of indexes) {
            await activateGPU( xordataset[j].inputs );
            await propagateGPU( xordataset[j].outputs, LEARNING_RATE);
        }

        if (epoch % 100 === 0) {
            let cost = 0;
            for (let j = 0; j < xordataset.length; j++) {
                let o = await activateGPU(  xordataset[j].inputs );
                for (let b=0; b<xordataset[j].outputs.length; b++)
                {
                    cost += Math.pow( xordataset[j].outputs[b] - o[b], 2);
                }
            }
            cost /= 4;
            console.log(`epoch ${epoch} mean squared error: ${cost}`);
        }
    }

    for (let i = 0; i < xordataset.length; i++)
    {
        const result = await activateGPU( xordataset[i].inputs );

        console.log(`for input ${xordataset[i].inputs} expected ${xordataset[i].outputs} predicted ${result[0].toFixed(4)} which is ${Math.round(result[0]) === xordataset[i].outputs[0] ? "correct" : "incorrect"}`);
    }



      copyBuffer( weightsBuffer, weights );
      copyBuffer( biasesBuffer,  biases );

}

//---------------------------------------------------------------------------

async function copyBuffer(gpuBuff, cpuBuff)
{
  /*
  Note: - make sure you add the 'GPUBufferUsage.COPY_SRC' flag to the buffer creation (copy the buffer data back)
  */
  // Retrieve results
  const readBuffer = device.createBuffer({size: gpuBuff.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });

  const commandEncoder = device.createCommandEncoder();
  commandEncoder.copyBufferToBuffer(gpuBuff, 0, readBuffer, 0, gpuBuff.size);
  device.queue.submit([commandEncoder.finish()]);

  await readBuffer.mapAsync(GPUMapMode.READ);
  const res = Array.from( new Float32Array(readBuffer.getMappedRange()) );
  cpuBuff.length = res.length;
  cpuBuff.forEach( (_,i)=>{cpuBuff[i]=res[i]; } );
  readBuffer.unmap();
}

//---------------------------------------------------------------------------

async function dumpBuffer(buff, str='data:')
{
  /*
  Note: - make sure you add the 'GPUBufferUsage.COPY_SRC' flag to the buffer creation (copy the buffer data back)
  */
  // Retrieve results
  const readBuffer = device.createBuffer({size: buff.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });

  const commandEncoder = device.createCommandEncoder();
  commandEncoder.copyBufferToBuffer(buff, 0, readBuffer, 0, buff.size);
  device.queue.submit([commandEncoder.finish()]);

  await readBuffer.mapAsync(GPUMapMode.READ);
  const results = Array.from( new Float32Array(readBuffer.getMappedRange()) );
  readBuffer.unmap();

  console.log( str, results );
}

//---------------------------------------------------------------------------

/*
but( 'load config', loadConfig );
but( 'save config', saveConfig );
but( 'reset weights and biases', initializeWeightsandBiases );
but( 'train JS network', trainJSNetwork );
but(' output JS network', outputJSNetwork );
but( 'train GPU network', trainGPUNetwork );
but(' output GPU network', outputGPUNetwork );
*/

//---------------------------------------------------------------------------

//--------------------------------------------------------------------------------------

xinitialize = ( args = { layers:[1,2,1], build:'cpu', iterations:10, learningrate:0.2 } ) => {

    if ( args.layers       ) xlayers       = args.layers;
    if ( args.build        ) xbuild        = args.build;
    if ( args.iterations   ) xiterations   = args.iterations;
    if ( args.learningrate ) xlearningrate = args.learningrate;

    if ( args.learningrate ) LEARNING_RATE = args.learningrate;
    if ( args.layers       ) layers = args.layers;

    maxLayerSize = [...layers].sort( (a,b)=>b-a )[0];

    weights  = Array( layers.length * maxLayerSize * maxLayerSize ).fill( 0 );
    biases   = Array( layers.length * maxLayerSize ).fill(0);
    loutputs = Array( layers.length * maxLayerSize ).fill(0);
    errors   = Array( layers.length * maxLayerSize ).fill(0);

    MAX_NEURONS_PER_LAYER = maxLayerSize;
    NUM_LAYERS            = layers.length;

    initializeWeightsandBiases();

    stConfig = {
      MAX_NEURONS_PER_LAYER : maxLayerSize,  // u32,
      NUM_LAYERS            : NUM_LAYERS,    // u32,
      LEARNING_RATE         : LEARNING_RATE*1000000, // f32, // 0.1, 0.2, ...
      ACTIVATION_TYPE       : 0              // u32  // enum - 0-sig, 1-relu, 2-reluleaky, ...
    }

    configBuffer    = createBuffer(new Uint32Array( Object.values(stConfig) ), GPUBufferUsage.STORAGE );
    layersBuffer    = createBuffer(new Uint32Array( layers ), GPUBufferUsage.STORAGE );
    inputsBuffer    = createBuffer(new Float32Array( layers[0]  ), GPUBufferUsage.STORAGE );
    resultsBuffer   = createBuffer(new Float32Array( layers[layers.length-1] ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC );
    expectedBuffer  = createBuffer(new Float32Array( layers[layers.length-1] ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC );
    weightsBuffer   = createBuffer(new Float32Array( weights.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
    biasesBuffer    = createBuffer(new Float32Array( biases.flat(4)  ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
    outputsBuffer   = createBuffer(new Float32Array( loutputs.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
    errorsBuffer    = createBuffer(new Float32Array( errors.flat(4) ), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);


}

xactivate = async ( ipp ) => {
    if ( xbuild == 'cpu' ) {
          return activate( ipp );
    }
    else if ( xbuild == 'gpu' ) {
          return await activateGPU( ipp );
    }
    console.assert('unknown build');
};

xpropagate = async ( op ) => {
    if ( xbuild == 'cpu' ) {
          return propagate( op );
    }
      if ( xbuild == 'gpu' ) {
          return propagateGPU( op );
    }
    console.assert('unknown build');
};

//--------------------------------------------------------------------------------------

Complete Code

Resources and Links