Skip to content

Commit

Permalink
Add memory function
Browse files Browse the repository at this point in the history
  • Loading branch information
@voidvoxel
voidvoxel committed Jun 16, 2024
1 parent 7371a23 commit 2a7840a
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 38 deletions.
41 changes: 33 additions & 8 deletions src/neural-network-gpu.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,18 +22,37 @@ import {
NeuralNetwork,
} from './neural-network';
import { release } from './utilities/kernel';
import { LossFunction, LossFunctionInputs, LossFunctionState } from './utilities/loss';
import { LossFunction, LossFunctionInputs, MemoryFunction, NeuralNetworkMemory } from './utilities/loss';

function loss(
this: IKernelFunctionThis,
actual: number,
expected: number,
inputs: LossFunctionInputs,
state: LossFunctionState
memory: NeuralNetworkMemory
) {
return expected - actual;
}

function updateMemory(
this: IKernelFunctionThis,
actual: number,
expected: number,
inputs: LossFunctionInputs,
memory: NeuralNetworkMemory,
memorySize: number,
loss: number
) {
const layer = this.thread.z;
const neuron = this.thread.y;
const signal = this.thread.x;

// Maintain the same signal magnitude.
return memory[layer][neuron][signal];
}

const DEFAULT_LOSS_FUNCTION = loss;
const DEFAULT_MEMORY_FUNCTION = updateMemory;

export interface INeuralNetworkGPUDatumFormatted {
input: KernelOutput;
Expand Down Expand Up @@ -402,7 +421,7 @@ export class NeuralNetworkGPU<
};

buildCalculateDeltas(): void {
let calcDeltas: GPUFunction<[number, number, LossFunctionInputs, LossFunctionState]>;
let calcDeltas: GPUFunction<[number, number, LossFunctionInputs, NeuralNetworkMemory]>;
switch (this.trainOpts.activation) {
case 'sigmoid':
calcDeltas = calcDeltasSigmoid;
Expand All @@ -423,6 +442,7 @@ export class NeuralNetworkGPU<
}

const loss: LossFunction = this._lossFunction ?? DEFAULT_LOSS_FUNCTION;
const updateMemory: MemoryFunction = this._memoryFunction ?? DEFAULT_MEMORY_FUNCTION;

calcDeltas = alias(
utils.getMinifySafeName(() => calcDeltas),
Expand All @@ -436,14 +456,14 @@ export class NeuralNetworkGPU<
// @ts-expect-error
this.backwardPropagate[this.outputLayer] = this.gpu.createKernelMap(
{
error: calcErrorOutput,
error: calcErrorOutput
},
function (
this: IKernelFunctionThis,
outputs: number[],
targets: number[],
inputs: LossFunctionInputs,
state: LossFunctionState
state: NeuralNetworkMemory
): number {
const output = outputs[this.thread.x];
const target = targets[this.thread.x];
Expand Down Expand Up @@ -503,9 +523,8 @@ export class NeuralNetworkGPU<

let output;
if (layer === this.outputLayer) {
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-expect-error
output = this.backwardPropagate[layer](this.outputs[layer], target, this.outputs[0], this.lossState);
// @ts-ignore
output = this.backwardPropagate[layer](this.outputs[layer], target, this.outputs[0], this.memory);
} else {
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-expect-error
Expand Down Expand Up @@ -731,11 +750,17 @@ export class NeuralNetworkGPU<
: (layerBiases as Float32Array)
)
);
const jsonLayerMemory = this.memory.map((layerMemory, layerIndex) =>
layerMemory.map(nodeMemory =>
Array.from(nodeMemory)
)
);
const jsonLayers: IJSONLayer[] = [];
for (let i = 0; i <= this.outputLayer; i++) {
jsonLayers.push({
weights: jsonLayerWeights[i] ?? [],
biases: jsonLayerBiases[i] ?? [],
memory: jsonLayerMemory[i] ?? []
});
}
return {
Expand Down
71 changes: 43 additions & 28 deletions src/neural-network.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ import { max } from './utilities/max';
import { mse } from './utilities/mse';
import { randos } from './utilities/randos';
import { zeros } from './utilities/zeros';
import { LossFunction, LossFunctionInputs, LossFunctionState } from './utilities/loss';
import { LossFunction, LossFunctionInputs, MemoryFunction, NeuralNetworkMemory } from './utilities/loss';
type NeuralNetworkFormatter =
| ((v: INumberHash) => Float32Array)
| ((v: number[]) => Float32Array);
Expand Down Expand Up @@ -44,7 +44,7 @@ function loss(
actual: number,
expected: number,
inputs: LossFunctionInputs,
state: LossFunctionState
state: NeuralNetworkMemory
) {
return expected - actual;
}
Expand All @@ -58,6 +58,7 @@ export type NeuralNetworkActivation =
export interface IJSONLayer {
biases: number[];
weights: number[][];
memory: number[][];
}

export interface INeuralNetworkJSON {
Expand All @@ -77,15 +78,15 @@ export interface INeuralNetworkOptions {
outputSize: number;
binaryThresh: number;
hiddenLayers?: number[];
lossStateSize: number;
memorySize: number;
}

export function defaults(): INeuralNetworkOptions {
return {
inputSize: 0,
outputSize: 0,
binaryThresh: 0.5,
lossStateSize: 1
memorySize: 1
};
}

Expand Down Expand Up @@ -119,8 +120,8 @@ export interface INeuralNetworkTrainOptions {
log: boolean | ((status: INeuralNetworkState) => void);
logPeriod: number;
loss?: LossFunction;
lossState?: LossFunctionState;
lossStateSize: number;
memory?: MemoryFunction;
memorySize: number;
leakyReluAlpha: number;
learningRate: number;
momentum: number;
Expand All @@ -141,7 +142,7 @@ export function trainDefaults(): INeuralNetworkTrainOptions {
log: false, // true to use console.log, when a function is supplied it is used
logPeriod: 10, // iterations between logging out
loss,
lossStateSize: 1,
memorySize: 1,
leakyReluAlpha: 0.01,
learningRate: 0.3, // multiply's against the input and the delta then adds to momentum
momentum: 0.1, // multiply's against the specified "change" then adds to learning rate for change
Expand Down Expand Up @@ -192,7 +193,7 @@ export class NeuralNetwork<
_formatInput: NeuralNetworkFormatter | null = null;
_formatOutput: NeuralNetworkFormatter | null = null;

_lossState: LossFunctionState;
_memory: NeuralNetworkMemory;

runInput: (input: Float32Array) => Float32Array = (input: Float32Array) => {
this.setActivation();
Expand All @@ -208,6 +209,7 @@ export class NeuralNetwork<
};

_lossFunction?: LossFunction;
_memoryFunction?: MemoryFunction;

// adam
biasChangesLow: Float32Array[] = [];
Expand All @@ -227,8 +229,9 @@ export class NeuralNetwork<
this.sizes = [inputSize].concat(hiddenLayers ?? []).concat([outputSize]);
}

const { lossStateSize } = this.options ?? 0;
this._lossState = this.trainOpts.lossState ?? this.replaceLossState(lossStateSize);
// Initialize memory matrix
const { memorySize } = this.options ?? 0;
this._memory = this.replaceMemory(memorySize);
}

/**
Expand Down Expand Up @@ -305,8 +308,8 @@ export class NeuralNetwork<
return this.sizes.length > 0;
}

public get lossState(): LossFunctionState {
return this._lossState;
public get memory(): NeuralNetworkMemory {
return this._memory;
}

run(input: Partial<InputType>): OutputType {
Expand Down Expand Up @@ -772,7 +775,7 @@ export class NeuralNetwork<
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);
error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.memory);
}
else error = target[node] - output;
} else {
Expand Down Expand Up @@ -805,7 +808,7 @@ export class NeuralNetwork<
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);
error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.memory);
}
else error = target[node] - output;
} else {
Expand Down Expand Up @@ -838,7 +841,7 @@ export class NeuralNetwork<
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);
error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.memory);
}
else error = target[node] - output;
} else {
Expand Down Expand Up @@ -870,7 +873,7 @@ export class NeuralNetwork<
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);
error = this._lossFunction.call(kernelFunctionThis, output, target[node], input, this.memory);
}
else error = target[node] - output;
} else {
Expand Down Expand Up @@ -1231,12 +1234,18 @@ export class NeuralNetwork<
const jsonLayerBiases = this.biases.map((layerBiases) =>
Array.from(layerBiases)
);
const jsonLayerMemory = this.memory.map(layerMemory =>
layerMemory.map(
nodeMemory => Array.from(nodeMemory)
)
);
const jsonLayers: IJSONLayer[] = [];
const outputLength = this.sizes.length - 1;
for (let i = 0; i <= outputLength; i++) {
jsonLayers.push({
weights: jsonLayerWeights[i] ?? [],
biases: jsonLayerBiases[i] ?? [],
memory: jsonLayerMemory[i] ?? []
});
}
return {
Expand Down Expand Up @@ -1281,9 +1290,15 @@ export class NeuralNetwork<
const layerBiases = this.biases.map((layerBiases, layerIndex) =>
Float32Array.from(jsonLayers[layerIndex].biases)
);
const layerMemory = this.memory.map((memory, layerIndex) =>
Array.from(jsonLayers[layerIndex].memory).map(nodeMemory =>
Float32Array.from(nodeMemory)
)
);
for (let i = 0; i <= this.outputLayer; i++) {
this.weights[i] = layerWeights[i] || [];
this.biases[i] = layerBiases[i] || [];
this.memory[i] = layerMemory[i] || [];
}
return this;
}
Expand Down Expand Up @@ -1387,23 +1402,23 @@ export class NeuralNetwork<
) => OutputType;
}

private createLossState(
lossStateSize: number
): LossFunctionState {
const lossState: LossFunctionState = [];
private createMemory(
memorySize: number
): NeuralNetworkMemory {
const memory: NeuralNetworkMemory = [];
for (let layer = 0; layer < this.sizes.length; layer++) {
lossState[layer] = [];
memory[layer] = [];
for (let neuron = 0; neuron < this.sizes.length; neuron++) {
lossState[layer][neuron] = new Float32Array(lossStateSize);
memory[layer][neuron] = new Float32Array(memorySize);
}
}
return lossState;
return memory;
}

private replaceLossState(
lossState: number | LossFunctionState
): LossFunctionState {
if (typeof lossState === "number") return this._lossState = this.createLossState(lossState);
return this._lossState = lossState;
private replaceMemory(
memory: number | NeuralNetworkMemory
): NeuralNetworkMemory {
if (typeof memory === "number") return this._memory = this.createMemory(memory);
return this._memory = memory;
}
}
14 changes: 12 additions & 2 deletions src/utilities/loss.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,22 @@ import { IKernelFunctionThis } from "gpu.js";

export type LossFunctionInputs = number[] | number[][] | number[][][] | Float32Array | Float32Array[] | Float32Array[][];

export type LossFunctionState = number[][][] | Float32Array[][];
export type NeuralNetworkMemory = Float32Array[][];

export type LossFunction = (
this: IKernelFunctionThis,
actual: number,
expected: number,
inputs: LossFunctionInputs,
state: LossFunctionState
memory: NeuralNetworkMemory
) => number;

export type MemoryFunction = (
this: IKernelFunctionThis,
actual: number,
expected: number,
inputs: LossFunctionInputs,
memory: NeuralNetworkMemory,
memorySize: number,
loss: number
) => number;

0 comments on commit 2a7840a

Please sign in to comment.