chore: clean leftover comments and println + fix shuffle (#69)

* remove print from flatten

* use stored gates

* ignore 0

* start only when cost exists

* generate wasm

* fix shuffle

* update iris example

Author: retraigo

crates/core/src/cpu/backend.rs

@@ -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;
                 }

crates/core/src/cpu/layers/embedding.rs

@@ -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};
 
@@ -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,
@@ -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),
+            ),
         }
     }
 
@@ -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()
     }
 }

crates/core/src/cpu/layers/flatten.rs

@@ -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()
     }
 

crates/core/src/cpu/layers/lstm.rs

@@ -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
@@ -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);
@@ -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)),
         }
     }
@@ -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)
@@ -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))
@@ -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> {
@@ -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));
 
@@ -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()));
@@ -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));
 
@@ -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()));
@@ -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 {

examples/classification/iris.ts

@@ -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

packages/core/src/backends/wasm/lib/netsaur.generated.js

@@ -4,7 +4,7 @@
 // deno-fmt-ignore-file
 /// <reference types="./netsaur.generated.d.ts" />
 
-// source-hash: 60e03f192e01f72854ad537d5cc565c02235b170
+// source-hash: 6fae1da844014c6af84634327825fa20fac7ef5f
 let wasm;
 let cachedInt32Memory0;
 
@@ -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) {