Generalize loss function for both CPU and GPU

Allow for using the same loss function between both CPU and GPU neural networks
voidvoxel · Jun 16, 2024 · 7371a23 · 7371a23
1 parent ba12f82
commit 7371a23
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/src/neural-network.ts b/src/neural-network.ts
@@ -768,7 +768,12 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") {
+            const kernelFunctionThis = { thread: { x: node, y: layer, z: 0 } };
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-ignore
+            error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.lossState);
+          }
           else error = target[node] - output;
         } else {
           const deltas = this.deltas[layer + 1];
@@ -796,7 +801,12 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") {
+            const kernelFunctionThis = { thread: { x: node, y: layer, z: 0 } };
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-ignore
+            error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.lossState);
+          }
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {
@@ -824,7 +834,12 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") {
+            const kernelFunctionThis = { thread: { x: node, y: layer, z: 0 } };
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-ignore
+            error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.lossState);
+          }
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {
@@ -851,7 +866,12 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") {
+            const kernelFunctionThis = { thread: { x: node, y: layer, z: 0 } };
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-ignore
+            error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.lossState);
+          }
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {

diff --git a/src/utilities/loss.ts b/src/utilities/loss.ts
@@ -1,8 +1,11 @@
+import { IKernelFunctionThis } from "gpu.js";
+
 export type LossFunctionInputs = number[] | number[][] | number[][][] | Float32Array | Float32Array[] | Float32Array[][];
 
 export type LossFunctionState = number[][][] | Float32Array[][];
 
 export type LossFunction = (
+  this: IKernelFunctionThis,
   actual: number,
   expected: number,
   inputs: LossFunctionInputs,