Skip to content

Commit

Permalink
chore: clean leftover comments and println + fix shuffle (#69)
Browse files Browse the repository at this point in the history 
* remove print from flatten

* use stored gates

* ignore 0

* start only when cost exists

* generate wasm

* fix shuffle

* update iris example
  • Loading branch information
@retraigo
retraigo authored Sep 30, 2024
1 parent 7248029 commit c1efc3e
Show file tree
Hide file tree
Showing 9 changed files with 167 additions and 121 deletions.
2 changes: 1 addition & 1 deletion crates/core/src/cpu/backend.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,7 @@ impl Backend {
total = 0.0;
}
}
if self.patience != 0 {
if self.patience != 0 && cost != 0.0 {
if best_cost < 0.0 {
best_cost = cost;
}
Expand Down
57 changes: 44 additions & 13 deletions crates/core/src/cpu/layers/embedding.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::ops::AddAssign;
use ndarray::{Array2, ArrayD, Axis, Ix2, IxDyn};
use std::ops::AddAssign;

use crate::{CPUInit, CPURegularizer, EmbeddingLayer, Init};

Expand All @@ -21,7 +21,10 @@ impl EmbeddingCPULayer {
pub fn new(config: EmbeddingLayer, size: IxDyn) -> Self {
let init = CPUInit::from(Init::Uniform);
let output_size = vec![size[0], size[1], config.embedding_size];
let embeddings = init.init(IxDyn(&[config.vocab_size, config.embedding_size]), 0, 0).into_dimensionality::<Ix2>().unwrap();
let embeddings = init
.init(IxDyn(&[config.vocab_size, config.embedding_size]), 0, 0)
.into_dimensionality::<Ix2>()
.unwrap();
let d_embeddings = Array2::zeros((config.vocab_size, config.embedding_size));
Self {
input_size: size,
Expand All @@ -32,7 +35,10 @@ impl EmbeddingCPULayer {
embeddings,
d_embeddings,
l_embeddings: Array2::zeros((config.vocab_size, config.embedding_size)),
regularizer: CPURegularizer::from(config.c.unwrap_or(0.0), config.l1_ratio.unwrap_or(1.0))
regularizer: CPURegularizer::from(
config.c.unwrap_or(0.0),
config.l1_ratio.unwrap_or(1.0),
),
}
}

Expand All @@ -48,21 +54,46 @@ impl EmbeddingCPULayer {
let input_indices: Vec<usize> = inputs.iter().map(|&x| x as usize).collect();
self.input_indices = input_indices.clone();
let embeddings = self.embeddings.select(Axis(0), input_indices.as_slice());
// let output_size = IxDyn(&self.output_size);
embeddings.into_shape_with_order(IxDyn(&[inputs.shape()[0], inputs.shape()[1], self.embedding_size])).unwrap()
// let output_size = IxDyn(&self.output_size);
embeddings
.into_shape_with_order(IxDyn(&[
inputs.shape()[0],
inputs.shape()[1],
self.embedding_size,
]))
.unwrap()
}

pub fn backward_propagate(&mut self, d_outputs: ArrayD<f32>) -> ArrayD<f32> {
let indices = Array2::from_shape_vec(Ix2(d_outputs.shape()[0], self.input_size[1]), self.input_indices.clone());
self.d_embeddings = Array2::zeros((self.d_embeddings.shape()[0], self.d_embeddings.shape()[1]));
d_outputs.axis_iter(Axis(0)).zip(indices).for_each(|(rec, i)| {
rec.axis_iter(Axis(0)).zip(i).for_each(|(grad, idx)| {
self.d_embeddings.index_axis_mut(Axis(0), idx).add_assign(&grad);
let indices = Array2::from_shape_vec(
Ix2(d_outputs.shape()[0], self.input_size[1]),
self.input_indices.clone(),
)
.unwrap();
self.d_embeddings.fill(0.0);
d_outputs
.axis_iter(Axis(0))
.zip(indices.axis_iter(Axis(0)))
.for_each(|(rec, i)| {
rec.axis_iter(Axis(0)).zip(i).for_each(|(grad, idx)| {
if idx != &0 {
self.d_embeddings
.index_axis_mut(Axis(0), *idx)
.add_assign(&grad);
}
});
});
});
self.l_embeddings = self.regularizer.coeff(&self.embeddings.clone().into_dyn()).into_dimensionality::<Ix2>().unwrap();
self.l_embeddings = self
.regularizer
.coeff(&self.embeddings.clone().into_dyn())
.into_dimensionality::<Ix2>()
.unwrap();
let mut input_size = self.input_size.clone();
input_size[0] = d_outputs.shape()[0];
ArrayD::from_shape_vec(input_size, self.input_indices.iter().map(|x| *x as f32).collect()).unwrap()
ArrayD::from_shape_vec(
input_size,
self.input_indices.iter().map(|x| *x as f32).collect(),
)
.unwrap()
}
}
1 change: 0 additions & 1 deletion crates/core/src/cpu/layers/flatten.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ impl FlattenCPULayer {

pub fn forward_propagate(&mut self, inputs: ArrayD<f32>) -> ArrayD<f32> {
let output_size = IxDyn(&[inputs.shape()[0], self.output_size[1]]);
println!("O {:?} {:?}", inputs.shape(), self.output_size);
inputs.into_shape_with_order(output_size).unwrap()
}

Expand Down
125 changes: 63 additions & 62 deletions crates/core/src/cpu/layers/lstm.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use crate::{CPUActivation, Activation, CPUInit, CPURegularizer, Init, LSTMLayer, Tensors, LayerNorm};
use core::f32;
use ndarray::{
concatenate, s, Array2, Array3, ArrayD, Axis, Dimension, Ix2, Ix3, IxDyn,
use crate::{
Activation, CPUActivation, CPUInit, CPURegularizer, Init, LSTMLayer, LayerNorm, Tensors,
};
use core::f32;
use ndarray::{concatenate, s, Array2, Array3, ArrayD, Axis, Dimension, Ix2, Ix3, IxDyn};
use std::ops::AddAssign;
/// Indices
/// 0 - Input Gate
Expand Down Expand Up @@ -30,12 +30,18 @@ pub struct LSTMCPULayer {
pub l_w_hh: Array3<f32>,
pub l_biases: Array2<f32>,

pub h: Array3<f32>,
pub c: Array3<f32>,
pub i_t: Array3<f32>,
pub f_t: Array3<f32>,
pub o_t: Array3<f32>,
pub g_t: Array3<f32>,

pub h_prev: Array2<f32>,
pub c_prev: Array2<f32>,

pub regularizer: CPURegularizer,
}

#[allow(unused_mut)]
impl LSTMCPULayer {
pub fn new(config: LSTMLayer, size: IxDyn, _tensors: Option<Tensors>) -> Self {
let return_sequences = config.return_sequences.unwrap_or(false);
Expand Down Expand Up @@ -68,14 +74,20 @@ impl LSTMCPULayer {
l_w_ih: Array3::zeros(weight_size),
l_w_hh: Array3::zeros((4, config.size, config.size)),
l_biases: Array2::zeros((4, config.size)),
h: Array3::zeros((size[1], size[0], config.size)),
c: Array3::zeros((size[1], size[0], config.size)),
i_t: Array3::zeros((size[1], size[0], config.size)),
f_t: Array3::zeros((size[1], size[0], config.size)),
o_t: Array3::zeros((size[1], size[0], config.size)),
g_t: Array3::zeros((size[1], size[0], config.size)),
h_prev: Array2::zeros((size[0], config.size)),
c_prev: Array2::zeros((size[0], config.size)),
regularizer: CPURegularizer::from(
config.c.unwrap_or(0.0),
config.l1_ratio.unwrap_or(1.0),
),

activation_h: CPUActivation::from(config.recurrent_activation.unwrap_or(Activation::Sigmoid)),
activation_h: CPUActivation::from(
config.recurrent_activation.unwrap_or(Activation::Sigmoid),
),
activation_o: CPUActivation::from(config.activation.unwrap_or(Activation::Tanh)),
}
}
Expand All @@ -92,10 +104,13 @@ impl LSTMCPULayer {
pub fn forward_propagate(&mut self, inputs: ArrayD<f32>) -> ArrayD<f32> {
self.inputs = inputs.into_dimensionality::<Ix3>().unwrap();
let output_size = self.w_ih.shape()[2];
self.h = Array3::zeros((self.inputs.shape()[1], self.inputs.shape()[0], output_size));
self.c = Array3::zeros((self.inputs.shape()[1], self.inputs.shape()[0], output_size));
let mut h_t = self.h.index_axis(Axis(0), 0).clone().to_owned();
let mut c_t = self.c.index_axis(Axis(0), 0).clone().to_owned();
let mut h_t = Array2::zeros((self.inputs.shape()[0], output_size));
let mut c_t = Array2::zeros(h_t.raw_dim());

self.i_t = Array3::zeros((self.inputs.shape()[1], self.inputs.shape()[0], output_size));
self.f_t = Array3::zeros(self.i_t.raw_dim());
self.o_t = Array3::zeros(self.i_t.raw_dim());
self.g_t = Array3::zeros(self.i_t.raw_dim());

let mut outputs = Array3::zeros(if self.return_sequences {
(self.inputs.shape()[0], self.inputs.shape()[1], output_size)
Expand All @@ -109,7 +124,8 @@ impl LSTMCPULayer {
.slice(s![.., t, ..])
.to_owned()
.into_dimensionality::<Ix2>()
.unwrap(); // Current input
.unwrap();

let i_t = (&x_t.dot(&self.w_ih.index_axis(Axis(0), 0))
+ &h_t.dot(&self.w_hh.index_axis(Axis(0), 0))
+ &self.biases.index_axis(Axis(0), 0))
Expand All @@ -127,56 +143,40 @@ impl LSTMCPULayer {
+ &self.biases.index_axis(Axis(0), 3))
.mapv(|x| (self.activation_o.activate)(&x));

self.i_t.index_axis_mut(Axis(0), t).assign(&i_t);
self.f_t.index_axis_mut(Axis(0), t).assign(&f_t);
self.o_t.index_axis_mut(Axis(0), t).assign(&o_t);
self.g_t.index_axis_mut(Axis(0), t).assign(&g_t);

c_t = &(&c_t * &f_t) + &(&g_t * &i_t);
h_t = &c_t.mapv(|x| (self.activation_o.activate)(&x)) * &o_t;

self.h.index_axis_mut(Axis(0), t).assign(&h_t);
self.c.index_axis_mut(Axis(0), t).assign(&c_t);

if self.return_sequences {
outputs.slice_mut(s![.., t, ..]).assign(&h_t);
}
}
self.h_prev = h_t.clone();
self.c_prev = c_t.clone();

if self.return_sequences {
outputs.into_dyn()
}
else {
} else {
h_t.into_dyn()
}
}

fn split_gates(
&self,
z: &Array2<f32>,
hidden_size: usize,
) -> (Array2<f32>, Array2<f32>, Array2<f32>, Array2<f32>) {
let i_t = z
.slice(ndarray::s![.., ..hidden_size])
.mapv(|x| (self.activation_h.activate)(&x));
let f_t = z
.slice(ndarray::s![.., hidden_size..2 * hidden_size])
.mapv(|x| (self.activation_h.activate)(&x));
let o_t = z
.slice(ndarray::s![.., 2 * hidden_size..3 * hidden_size])
.mapv(|x| (self.activation_h.activate)(&x));
let g_t = z
.slice(ndarray::s![.., 3 * hidden_size..])
.mapv(|x| (self.activation_o.activate)(&x));

(i_t, f_t, o_t, g_t)
}
pub fn backward_propagate(&mut self, d_outputs: ArrayD<f32>) -> ArrayD<f32> {
match d_outputs.shape().len() {
2 => {
let d_inputs = self.backward_propagate_2d(d_outputs.into_dimensionality::<Ix2>().unwrap());
let d_inputs =
self.backward_propagate_2d(d_outputs.into_dimensionality::<Ix2>().unwrap());
d_inputs.into_dyn()
}
3 => {
let d_inputs = self.backward_propagate_3d(d_outputs.into_dimensionality::<Ix3>().unwrap());
let d_inputs =
self.backward_propagate_3d(d_outputs.into_dimensionality::<Ix3>().unwrap());
d_inputs.into_dyn()
}
_ => d_outputs
_ => d_outputs,
}
}
pub fn backward_propagate_3d(&mut self, d_outputs: Array3<f32>) -> Array3<f32> {
Expand All @@ -189,8 +189,8 @@ impl LSTMCPULayer {
self.d_w_hh = Array3::zeros((4, hidden_size, hidden_size));
self.d_biases = Array2::zeros((4, hidden_size));

let h_prev = self.h.index_axis(Axis(0), sequence_length - 1);
let c_prev = self.c.index_axis(Axis(0), sequence_length - 1);
let h_prev = self.h_prev.clone();
let c_prev = self.c_prev.clone();

let mut d_inputs = Array3::<f32>::zeros((batch_size, sequence_length, input_size));

Expand Down Expand Up @@ -224,27 +224,28 @@ impl LSTMCPULayer {
.into_dimensionality::<Ix2>()
.unwrap();

let d_h = d_outputs.slice(s![.., t, ..]).clone().to_owned().into_dimensionality::<Ix2>().unwrap();
let d_h = d_outputs
.slice(s![.., t, ..])
.clone()
.to_owned()
.into_dimensionality::<Ix2>()
.unwrap();

d_h_prev = d_h_prev + d_h;
// clip_gradients(&mut d_h_prev, 5f32);

let gates = x_t.dot(&w_ih)
+ h_prev.dot(&w_hh)
+ &self.biases.to_shape(4 * hidden_size).unwrap();
let (i_t, f_t, o_t, g_t) = self.split_gates(&gates, hidden_size);
let i_t = self.i_t.index_axis(Axis(0), t);
let f_t = self.f_t.index_axis(Axis(0), t);
let o_t = self.o_t.index_axis(Axis(0), t);
let g_t = self.g_t.index_axis(Axis(0), t);

let d_tanned_c = &d_h_prev * &o_t * c_prev.map(|x| (self.activation_o.activate)(&x));
let mut d_c_t = d_tanned_c + &d_c_prev;
// clip_gradients(&mut d_c_t, 5f32);

let d_o_t = &d_h_prev * &c_prev.mapv(|x| (self.activation_o.activate)(&x));
let d_f_t = &d_c_t * &c_prev * &f_t.map(|x| (self.activation_h.prime)(x));
let d_i_t = &d_c_t * &g_t * &i_t.map(|x| (self.activation_h.prime)(x));
let d_g_t = &d_c_t * &i_t * &g_t.map(|x| (self.activation_o.prime)(x));
let d_gates = concatenate![Axis(1), d_i_t, d_f_t, d_o_t, d_g_t];
// println!("ADD {:?}", concatenate![Axis(0), d_i_t, d_f_t, d_o_t, d_g_t].shape());
// println!("DG {:?} {:?}", d_gates.shape(), &w_ih.shape());
d_inputs
.slice_mut(s![.., t, ..])
.assign(&d_gates.dot(&w_ih.t()));
Expand Down Expand Up @@ -303,8 +304,8 @@ impl LSTMCPULayer {
self.d_w_hh = Array3::zeros((4, hidden_size, hidden_size));
self.d_biases = Array2::zeros((4, hidden_size));

let h_prev = self.h.index_axis(Axis(0), sequence_length - 1);
let c_prev = self.c.index_axis(Axis(0), sequence_length - 1);
let h_prev = self.h_prev.clone();
let c_prev = self.c_prev.clone();

let mut d_inputs = Array3::<f32>::zeros((batch_size, sequence_length, input_size));

Expand Down Expand Up @@ -339,22 +340,21 @@ impl LSTMCPULayer {
.into_dimensionality::<Ix2>()
.unwrap();

let gates = x_t.dot(&w_ih)
+ h_prev.dot(&w_hh)
+ &self.biases.to_shape(4 * hidden_size).unwrap();
let (i_t, f_t, o_t, g_t) = self.split_gates(&gates, hidden_size);
let i_t = self.i_t.index_axis(Axis(0), t);
let f_t = self.f_t.index_axis(Axis(0), t);
let o_t = self.o_t.index_axis(Axis(0), t);
let g_t = self.g_t.index_axis(Axis(0), t);

let d_tanned_c = &d_h_prev * &o_t * c_prev.map(|x| (self.activation_o.activate)(&x));
let mut d_c_t = d_tanned_c + &d_c_prev;
// clip_gradients(&mut d_c_t, 5f32);

let d_o_t = &d_h_prev * &c_prev.mapv(|x| (self.activation_o.activate)(&x));
let d_f_t = &d_c_t * &c_prev * &f_t.map(|x| (self.activation_h.prime)(x));
let d_i_t = &d_c_t * &g_t * &i_t.map(|x| (self.activation_h.prime)(x));
let d_g_t = &d_c_t * &i_t * &g_t.map(|x| (self.activation_o.prime)(x));

let d_gates = concatenate![Axis(1), d_i_t, d_f_t, d_o_t, d_g_t];
// println!("OT: {:?}\n\nFT {:?}", &d_h_prev, &d_c_prev);

d_inputs
.slice_mut(s![.., t, ..])
.assign(&d_gates.dot(&w_ih.t()));
Expand Down Expand Up @@ -405,6 +405,7 @@ impl LSTMCPULayer {
}
}

#[allow(dead_code)]
fn clip_gradients(grad: &mut Array2<f32>, threshold: f32) -> () {
let norm = grad.mapv(|x| x.powi(2)).sum().sqrt();
if norm > threshold {
Expand Down
8 changes: 6 additions & 2 deletions examples/classification/iris.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,12 @@ const net = new Sequential({

// Define each layer of the network
layers: [
// A dense layer with 16 neurons
DenseLayer({ size: [16] }),
// A dense layer with 8 neurons
DenseLayer({ size: [8] }),
// A ReLu activation layer
ReluLayer(),
// A dense layer with 8 neurons
DenseLayer({ size: [8] }),
// A ReLu activation layer
ReluLayer(),
// A dense layer with 3 neurons
Expand Down
6 changes: 3 additions & 3 deletions packages/core/src/backends/wasm/lib/netsaur.generated.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
// deno-fmt-ignore-file
/// <reference types="./netsaur.generated.d.ts" />

// source-hash: 60e03f192e01f72854ad537d5cc565c02235b170
// source-hash: 6fae1da844014c6af84634327825fa20fac7ef5f
let wasm;
let cachedInt32Memory0;

Expand Down Expand Up @@ -223,10 +223,10 @@ const imports = {
__wbindgen_object_drop_ref: function (arg0) {
takeObject(arg0);
},
__wbg_log_6f7dfa87fad40a57: function (arg0, arg1) {
__wbg_log_860f982bd7b08de8: function (arg0, arg1) {
console.log(getStringFromWasm0(arg0, arg1));
},
__wbg_now_de5fe0de473bcd7d: typeof Date.now == "function"
__wbg_now_aa5f80cbe756bd45: typeof Date.now == "function"
? Date.now
: notDefined("Date.now"),
__wbindgen_number_new: function (arg0) {
Expand Down
Binary file modified packages/core/src/backends/wasm/lib/netsaur_bg.wasm
View file
Binary file not shown.
Loading

0 comments on commit c1efc3e

Please sign in to comment.