New layer architecture

.gitignore

@@ -1,5 +1,7 @@
 # will have compiled files and executables
 target
 
+Cargo.lock
+
 # These are backup files generated by rustfmt
 **/*.rs.bk

Cargo.toml

@@ -2,7 +2,7 @@
 members = [
 		"greenglas", "coaster", "coaster-nn", "coaster-blas", "juice", "rust-blas",
 		"rcudnn/cudnn", "rcudnn/cudnn-sys", "rcublas/cublas", "rcublas/cublas-sys",
-		"juice-examples/juice-utils", "juice-examples/mackey-glass-rnn-regression",
+#		"juice-examples/juice-utils", "juice-examples/mackey-glass-rnn-regression",
 		"juice-examples/mnist-image-multiclass-classification"]
 
 exclude = [ "./rcudnn", "./rcublas", "./juice-examples"]

juice-examples/mnist-image-multiclass-classification/src/main.rs

@@ -10,8 +10,7 @@ use co::frameworks::cuda::get_cuda_backend;
 #[cfg(not(feature = "cuda"))]
 use co::frameworks::native::get_native_backend;
 use co::prelude::*;
-use juice::layer::*;
-use juice::layers::*;
+use juice::net::*;
 use juice::solver::*;
 use juice::util::*;
 use juice_utils::{download_datasets, unzip_datasets};
@@ -131,7 +130,7 @@ fn main() {
         );
     }
 }
-
+/*
 #[cfg(all(feature = "cuda"))]
 fn add_conv_net(
     mut net_cfg: SequentialConfig,
@@ -164,7 +163,7 @@ fn add_conv_net(
         "linear1",
         LinearConfig { output_size: 500 },
     ));
-    net_cfg.add_layer(LayerConfig::new("sigmoid", LayerType::Sigmoid));
+    net_cfg.add_layer(LayerConfig::new("sigmoid", LayerConfig::Sigmoid));
     net_cfg.add_layer(LayerConfig::new(
         "linear2",
         LinearConfig { output_size: 10 },
@@ -192,26 +191,25 @@ fn add_mlp(
 ) -> SequentialConfig {
     net_cfg.add_layer(LayerConfig::new(
         "reshape",
-        LayerType::Reshape(ReshapeConfig::of_shape(&[batch_size, pixel_count])),
+        LayerConfig::Reshape(ReshapeConfig::of_shape(&[batch_size, pixel_count])),
     ));
     net_cfg.add_layer(LayerConfig::new(
         "linear1",
-        LayerType::Linear(LinearConfig { output_size: 1568 }),
+        LayerConfig::Linear(LinearConfig { output_size: 1568 }),
     ));
-    net_cfg.add_layer(LayerConfig::new("sigmoid", LayerType::Sigmoid));
+    net_cfg.add_layer(LayerConfig::new("sigmoid", LayerConfig::Sigmoid));
     net_cfg.add_layer(LayerConfig::new(
         "linear2",
-        LayerType::Linear(LinearConfig { output_size: 10 }),
+        LayerConfig::Linear(LinearConfig { output_size: 10 }),
     ));
     net_cfg
 }
-
-fn add_linear_net(mut net_cfg: SequentialConfig) -> SequentialConfig {
-    net_cfg.add_layer(LayerConfig::new(
+*/
+fn add_linear_net(net_cfg: SequentialConfig) -> SequentialConfig{
+    net_cfg.with_layer(
         "linear",
-        LayerType::Linear(LinearConfig { output_size: 10 }),
-    ));
-    net_cfg
+        LayerConfig::Linear(LinearConfig { output_size: 10 }),
+    )
 }
 
 fn run_mnist(
@@ -250,25 +248,19 @@ fn run_mnist(
     let momentum = momentum.unwrap_or(0f32);
 
     let mut net_cfg = SequentialConfig::default();
-    net_cfg.add_input("data", &[batch_size, pixel_dim, pixel_dim]);
-    net_cfg.force_backward = true;
 
-    net_cfg = match &*model_name.unwrap_or("none".to_owned()) {
-        "conv" => add_conv_net(net_cfg, batch_size, pixel_dim),
-        "mlp" => add_mlp(net_cfg, batch_size, pixel_count),
-        "linear" => add_linear_net(net_cfg),
+    match &*model_name.unwrap_or("none".to_owned()) {
+        // "conv" => add_conv_net(net_cfg, batch_size, pixel_dim),
+        // "mlp" => add_mlp(net_cfg, batch_size, pixel_count),
+        "linear" => net_cfg = add_linear_net(net_cfg),
         _ => panic!("Unknown model. Try one of [linear, mlp, conv]"),
     };
 
-    net_cfg.add_layer(LayerConfig::new("log_softmax", LayerType::LogSoftmax));
+    net_cfg = net_cfg.with_layer("log_softmax", LayerConfig::LogSoftmax);
 
-    let mut classifier_cfg = SequentialConfig::default();
-    classifier_cfg.add_input("network_out", &[batch_size, 10]);
-    classifier_cfg.add_input("label", &[batch_size, 1]);
     // set up nll loss
-    let nll_layer_cfg = NegativeLogLikelihoodConfig { num_classes: 10 };
-    let nll_cfg = LayerConfig::new("nll", LayerType::NegativeLogLikelihood(nll_layer_cfg));
-    classifier_cfg.add_layer(nll_cfg);
+    let classifier_cfg =
+        LayerConfig::NegativeLogLikelihood(NegativeLogLikelihoodConfig { num_classes: 10 });
 
     // set up backends
     #[cfg(all(feature = "cuda"))]
@@ -283,9 +275,14 @@ fn run_mnist(
         momentum,
         ..SolverConfig::default()
     };
-    solver_cfg.network = LayerConfig::new("network", net_cfg);
-    solver_cfg.objective = LayerConfig::new("classifier", classifier_cfg);
-    let mut solver = Solver::from_config(backend.clone(), backend.clone(), &solver_cfg);
+    solver_cfg.network = LayerConfig::Sequential(net_cfg);
+    solver_cfg.objective = classifier_cfg;
+    let mut solver = Solver::from_config(
+        backend.clone(),
+        &solver_cfg,
+        &[vec![pixel_dim, pixel_dim]],
+        &vec![1],
+    );
 
     // set up confusion matrix
     let mut classification_evaluator = ::juice::solver::ConfusionMatrix::new(10);
@@ -311,9 +308,8 @@ fn run_mnist(
             targets.push(label_val as usize);
         }
         // train the network!
-        let infered_out = solver.train_minibatch(inp_lock.clone(), label_lock.clone());
+        let mut infered = solver.train_minibatch(inp_lock.clone(), label_lock.clone());
 
-        let mut infered = infered_out.write().unwrap();
         let predictions = classification_evaluator.get_predictions(&mut infered);
 
         classification_evaluator.add_samples(&predictions, &targets);

juice/build.rs

@@ -4,4 +4,10 @@ fn main() {
         .file("capnp/juice.capnp")
         .run()
         .expect("capnpc schema compiler command must succeed");
+
+    capnpc::CompilerCommand::new()
+        .src_prefix("capnp")
+        .file("capnp/net.capnp")
+        .run()
+        .expect("capnpc schema compiler command must succeed");
 }

juice/capnp/net.capnp

@@ -0,0 +1,34 @@
+# Types required for serialization of a Network.
+
+@0xd42882f7f90ce348;
+
+struct SequentialConfig {
+}
+
+struct LinearConfig {
+    outputSize @0 :UInt64;
+}
+
+struct NegativeLogLikelihoodConfig {
+}
+
+struct LayerConfig {
+    layerType :union {
+        sequential @0 :SequentialConfig; 
+        linear @1 :LinearConfig;
+        logSoftmax @2 :Void;
+        relu @3 :Void;
+        sigmoid @4 :Void;
+        meanSquaredError @5 :Void;
+        negativeLogLikelihood @6 :NegativeLogLikelihoodConfig;
+    }
+}
+
+struct TensorShape {
+    shape @0 :List(UInt64);
+}
+
+struct Network {
+    config @0 :LayerConfig;
+    inputs @1 :List(TensorShape);
+}

juice/examples/benchmarks.rs

@@ -1,3 +1,5 @@
+/*
+
 #[macro_use]
 extern crate timeit;
 
@@ -583,3 +585,7 @@ fn bench_vgg_a() {
         }
     }
 }
+*/
+
+
+fn main() {}

juice/src/lib.rs

@@ -117,10 +117,12 @@ extern crate coaster_blas as coblas;
 extern crate coaster_nn as conn;
 extern crate num;
 extern crate rand;
-pub mod layer;
-pub mod layers;
+//pub mod layer;
+//pub mod layers;
+pub mod net;
 pub mod solver;
 pub mod solvers;
+pub mod train;
 pub mod weight;
 
 mod capnp_util;
@@ -130,3 +132,6 @@ pub mod util;
 mod juice_capnp {
     include!(concat!(env!("OUT_DIR"), "/juice_capnp.rs"));
 }
+mod net_capnp {
+    include!(concat!(env!("OUT_DIR"), "/net_capnp.rs"));
+}

juice/src/net/activation/mod.rs

@@ -0,0 +1,5 @@
+mod relu;
+mod sigmoid;
+
+pub use relu::*;
+pub use sigmoid::*;

juice/src/net/activation/relu.rs

@@ -0,0 +1,53 @@
+use crate::co::IBackend;
+use crate::conn;
+use crate::net::{Context, Descriptor, Layer};
+
+#[derive(Debug, Clone)]
+pub struct Relu {
+    descriptor: Descriptor,
+}
+
+impl Relu {
+    pub fn new(mut descriptor: Descriptor) -> Self {
+        assert_eq!(descriptor.inputs().len(), 1); // Should only be one input.
+
+        descriptor.add_output(descriptor.input(0).unit_shape().clone());
+
+        Relu {
+            descriptor: descriptor,
+        }
+    }
+}
+
+impl<B: IBackend + conn::Relu<f32>> Layer<B> for Relu {
+    fn compute_output(&self, backend: &B, context: &mut Context) {
+        let input = context.get_data(self.descriptor.input(0));
+        let output = context.acquire_data(self.descriptor.output(0));
+        backend
+            .relu(&input.borrow(), &mut output.borrow_mut())
+            .unwrap();
+    }
+
+    fn compute_gradients(&self, backend: &B, context: &mut Context) {
+        let input = context.get_data(self.descriptor.input(0));
+        let output = context.get_data(self.descriptor.output(0));
+        let output_gradient = context.get_data_gradient(self.descriptor.output(0));
+        let input_gradient = context.acquire_data_gradient(self.descriptor.input(0));
+        backend
+            .relu_grad(
+                &output.borrow(),
+                &output_gradient.borrow(),
+                &input.borrow(),
+                &mut input_gradient.borrow_mut(),
+            )
+            .unwrap();
+    }
+
+    fn descriptor(&self) -> &Descriptor {
+        &self.descriptor
+    }
+
+    fn descriptor_mut(&mut self) -> &mut Descriptor {
+        &mut self.descriptor
+    }
+}

juice/src/net/activation/sigmoid.rs

@@ -0,0 +1,53 @@
+use crate::co::IBackend;
+use crate::conn;
+use crate::net::{Context, Descriptor, Layer};
+
+#[derive(Debug, Clone)]
+pub struct Sigmoid {
+    descriptor: Descriptor,
+}
+
+impl Sigmoid {
+    pub fn new(mut descriptor: Descriptor) -> Self {
+        assert_eq!(descriptor.inputs().len(), 1); // Should only be one input.
+
+        descriptor.add_output(descriptor.input(0).unit_shape().clone());
+
+        Sigmoid {
+            descriptor: descriptor,
+        }
+    }
+}
+
+impl<B: IBackend + conn::Sigmoid<f32>> Layer<B> for Sigmoid {
+    fn compute_output(&self, backend: &B, context: &mut Context) {
+        let input = context.get_data(self.descriptor.input(0));
+        let output = context.acquire_data(self.descriptor.output(0));
+        backend
+            .sigmoid(&input.borrow(), &mut output.borrow_mut())
+            .unwrap();
+    }
+
+    fn compute_gradients(&self, backend: &B, context: &mut Context) {
+        let input = context.get_data(self.descriptor.input(0));
+        let output = context.get_data(self.descriptor.output(0));
+        let output_gradient = context.get_data_gradient(self.descriptor.output(0));
+        let input_gradient = context.acquire_data_gradient(self.descriptor.input(0));
+        backend
+            .sigmoid_grad(
+                &output.borrow(),
+                &output_gradient.borrow(),
+                &input.borrow(),
+                &mut input_gradient.borrow_mut(),
+            )
+            .unwrap();
+    }
+
+    fn descriptor(&self) -> &Descriptor {
+        &self.descriptor
+    }
+
+    fn descriptor_mut(&mut self) -> &mut Descriptor {
+        &mut self.descriptor
+    }
+}