Design custom loss function for autoencoders

README.md

@@ -43,7 +43,7 @@ GPU accelerated Neural networks in JavaScript for Browsers and Node.js
     - [For training with NeuralNetwork](#for-training-with-neuralnetwork)
     - [For training with `RNNTimeStep`, `LSTMTimeStep` and `GRUTimeStep`](#for-training-with-rnntimestep-lstmtimestep-and-grutimestep)
     - [For training with `RNN`, `LSTM` and `GRU`](#for-training-with-rnn-lstm-and-gru)
-    - [For training with `AE`](#for-training-with-ae)
+    - [For training with `AutoencoderGPU`](#for-training-with-ae)
   - [Training Options](#training-options)
   - [Async Training](#async-training)
   - [Cross Validation](#cross-validation)
@@ -318,7 +318,7 @@ net.train([
 const output = net.run('I feel great about the world!'); // 'happy'
 ```
 
-#### For training with `AE`
+#### For training with `AutoencoderGPU`
 
 Each training pattern can either:
 
@@ -328,7 +328,7 @@ Each training pattern can either:
 Training an autoencoder to compress the values of a XOR calculation:
 
 ```javascript
-const net = new brain.AE(
+const net = new brain.AutoencoderGPU(
   {
     hiddenLayers: [ 5, 2, 5 ]
   }
@@ -644,7 +644,7 @@ The user interface used:
 
 - [`brain.NeuralNetwork`](src/neural-network.ts) - [Feedforward Neural Network](https://en.wikipedia.org/wiki/Feedforward_neural_network) with backpropagation
 - [`brain.NeuralNetworkGPU`](src/neural-network-gpu.ts) - [Feedforward Neural Network](https://en.wikipedia.org/wiki/Feedforward_neural_network) with backpropagation, GPU version
-- [`brain.AE`](src/autoencoder.ts) - [Autoencoder or "AE"](https://en.wikipedia.org/wiki/Autoencoder) with backpropogation and GPU support
+- [`brain.AutoencoderGPU`](src/autoencoder.ts) - [Autoencoder or "AutoencoderGPU"](https://en.wikipedia.org/wiki/Autoencoder) with backpropogation and GPU support
 - [`brain.recurrent.RNNTimeStep`](src/recurrent/rnn-time-step.ts) - [Time Step Recurrent Neural Network or "RNN"](https://en.wikipedia.org/wiki/Recurrent_neural_network)
 - [`brain.recurrent.LSTMTimeStep`](src/recurrent/lstm-time-step.ts) - [Time Step Long Short Term Memory Neural Network or "LSTM"](https://en.wikipedia.org/wiki/Long_short-term_memory)
 - [`brain.recurrent.GRUTimeStep`](src/recurrent/gru-time-step.ts) - [Time Step Gated Recurrent Unit or "GRU"](https://en.wikipedia.org/wiki/Gated_recurrent_unit)

src/autoencoder.test.ts

@@ -1,4 +1,4 @@
-import AE from "./autoencoder";
+import AutoencoderGPU from "./autoencoder";
 
 const trainingData = [
   [0, 0, 0],
@@ -7,9 +7,9 @@ const trainingData = [
   [1, 1, 0]
 ];
 
-const xornet = new AE<number[], number[]>(
+const xornet = new AutoencoderGPU<number[], number[]>(
   {
-    decodedSize: 3,
+    inputSize: 3,
     hiddenLayers: [ 5, 2, 5 ]
   }
 );

src/autoencoder.ts

@@ -1,41 +1,54 @@
-import { KernelOutput, Texture, TextureArrayOutput } from "gpu.js";
-import { IJSONLayer, INeuralNetworkData, INeuralNetworkDatum, INeuralNetworkTrainOptions } from "./neural-network";
+import { IKernelFunctionThis, KernelOutput, Texture, TextureArrayOutput } from "gpu.js";
+import { IJSONLayer, INeuralNetworkData, INeuralNetworkDatum, INeuralNetworkTrainOptions, NeuralNetworkIO, NeuralNetworkRAM } from "./neural-network";
 import { INeuralNetworkGPUOptions, NeuralNetworkGPU } from "./neural-network-gpu";
 import { INeuralNetworkState } from "./neural-network-types";
 import { UntrainedNeuralNetworkError } from "./errors/untrained-neural-network-error";
 
-export interface IAEOptions {
-  binaryThresh: number;
-  decodedSize: number;
-  hiddenLayers: number[];
+function loss(
+  this: IKernelFunctionThis,
+  actual: number,
+  expected: number,
+  inputs: NeuralNetworkIO,
+  ram: NeuralNetworkRAM
+) {
+  let error = expected - actual;
+
+  // if ( o ≈ i0 ) then return 10% of the loss value.
+  // Otherwise, return 1000% of the full loss value.
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  if (Math.round(actual) !== Math.round(inputs[this.thread.x])) error *= 32;
+  else error *= 0.03125;
+
+  return error;
 }
 
 /**
  * An autoencoder learns to compress input data down to relevant features and reconstruct input data from its compressed representation.
  */
-export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeuralNetworkData> {
+export class AutoencoderGPU<DecodedData extends INeuralNetworkData, EncodedData extends INeuralNetworkData> extends NeuralNetworkGPU<DecodedData, DecodedData> {
   private decoder?: NeuralNetworkGPU<EncodedData, DecodedData>;
-  private denoiser: NeuralNetworkGPU<DecodedData, DecodedData>;
 
   constructor (
-    options?: Partial<IAEOptions>
+    options?: Partial<INeuralNetworkGPUOptions>
   ) {
     // Create default options for the autoencoder.
     options ??= {};
 
-    // Create default options for the autoencoder's denoiser subnet.
-    const denoiserOptions: Partial<INeuralNetworkGPUOptions> = {};
-
     // Inherit the binary threshold of the parent autoencoder.
-    denoiserOptions.binaryThresh = options.binaryThresh;
+    options.binaryThresh = options.binaryThresh;
     // Inherit the hidden layers of the parent autoencoder.
-    denoiserOptions.hiddenLayers = options.hiddenLayers;
+    options.hiddenLayers = options.hiddenLayers;
+
+    const decodedSize = options.inputSize ?? options.outputSize ?? 1;
 
     // Define the denoiser subnet's input and output sizes.
-    if (options.decodedSize) denoiserOptions.inputSize = denoiserOptions.outputSize = options.decodedSize;
+    if (decodedSize) options.inputSize = options.outputSize = decodedSize;
+
+    options.loss ??= loss;
 
-    // Create the denoiser subnet of the autoencoder.
-    this.denoiser = new NeuralNetworkGPU<DecodedData, DecodedData>(options);
+    // Create the autoencoder.
+    super(options);
   }
 
   /**
@@ -51,7 +64,7 @@ export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeu
     // denoising and anomaly detection; there are other specialized topologies
     // better suited for these tasks anyways, many of which can be implemented
     // by using an autoencoder.
-    return this.denoiser.run(input);
+    return this.run(input);
   }
 
   /**
@@ -76,10 +89,10 @@ export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeu
    */
   encode(input: DecodedData): EncodedData {
     // If the decoder has not been trained yet, throw an error.
-    if (!this.denoiser) throw new UntrainedNeuralNetworkError(this);
+    if (!this) throw new UntrainedNeuralNetworkError(this);
 
     // Process the input.
-    this.denoiser.run(input);
+    this.run(input);
 
     // Get the auto-encoded input.
     let encodedInput: TextureArrayOutput = this.encodedLayer as TextureArrayOutput;
@@ -131,14 +144,15 @@ export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeu
    * @param {Partial<INeuralNetworkTrainOptions>} options
    * @returns {INeuralNetworkState}
    */
-  train(data: DecodedData[], options?: Partial<INeuralNetworkTrainOptions>): INeuralNetworkState {
+  train(data: Partial<DecodedData>[] | INeuralNetworkDatum<Partial<DecodedData>, Partial<DecodedData>>[], options?: Partial<INeuralNetworkTrainOptions>): INeuralNetworkState {
     const preprocessedData: INeuralNetworkDatum<Partial<DecodedData>, Partial<DecodedData>>[] = [];
 
+    if (data.length && data.length > 0)
     for (let datum of data) {
-      preprocessedData.push( { input: datum, output: datum } );
+      preprocessedData.push( { input: datum as Partial<DecodedData>, output: datum as Partial<DecodedData> } );
     }
 
-    const results = this.denoiser.train(preprocessedData, options);
+    const results = super.train(preprocessedData, options);
 
     this.decoder = this.createDecoder();
 
@@ -151,12 +165,12 @@ export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeu
    * @returns {NeuralNetworkGPU<EncodedData, DecodedData>}
    */
   private createDecoder() {
-    const json = this.denoiser.toJSON();
+    const json = this.toJSON();
 
     const layers: IJSONLayer[] = [];
     const sizes: number[] = [];
 
-    for (let i = this.encodedLayerIndex; i < this.denoiser.sizes.length; i++) {
+    for (let i = this.encodedLayerIndex; i < this.sizes.length; i++) {
       layers.push(json.layers[i]);
       sizes.push(json.sizes[i]);
     }
@@ -175,15 +189,15 @@ export class AE<DecodedData extends INeuralNetworkData, EncodedData extends INeu
    * Get the layer containing the encoded representation.
    */
   private get encodedLayer(): KernelOutput {
-    return this.denoiser.outputs[this.encodedLayerIndex];
+    return this.outputs[this.encodedLayerIndex];
   }
 
   /**
    * Get the offset of the encoded layer.
    */
   private get encodedLayerIndex(): number {
-    return Math.round(this.denoiser.outputs.length * 0.5) - 1;
+    return Math.round(this.outputs.length * 0.5) - 1;
   }
 }
 
-export default AE;
+export default AutoencoderGPU;

src/errors/untrained-neural-network-error.ts

@@ -1,6 +1,9 @@
+import { INeuralNetworkData, NeuralNetwork } from "../neural-network";
+import { NeuralNetworkGPU } from "../neural-network-gpu";
+
 export class UntrainedNeuralNetworkError extends Error {
   constructor (
-    neuralNetwork: any
+    neuralNetwork: object
   ) {
     super(`Cannot run a ${neuralNetwork.constructor.name} before it is trained.`);
   }

src/index.ts

@@ -1,5 +1,5 @@
 import * as activation from './activation';
-import { AE } from './autoencoder';
+import { AutoencoderGPU } from './autoencoder';
 import CrossValidate from './cross-validate';
 import { FeedForward } from './feed-forward';
 import * as layer from './layer';
@@ -54,7 +54,7 @@ const utilities = {
 
 export {
   activation,
-  AE,
+  AutoencoderGPU,
   CrossValidate,
   likely,
   layer,

src/neural-network.ts

@@ -23,7 +23,7 @@ export type LossFunction = (
   actual: number,
   expected: number,
   inputs: NeuralNetworkIO,
-  memory: NeuralNetworkRAM
+  ram: NeuralNetworkRAM
 ) => number;
 
 export type RAMFunction = (
@@ -108,7 +108,7 @@ function loss(
   actual: number,
   expected: number,
   inputs: NeuralNetworkIO,
-  memory: NeuralNetworkRAM
+  ram: NeuralNetworkRAM
 ) {
   return expected - actual;
 }