From 251e4d02d29538b2d3d3fa9ddccbc687ab5522b7 Mon Sep 17 00:00:00 2001 From: Adrien Burgun Date: Sun, 7 May 2023 11:28:27 +0200 Subject: [PATCH] :sparkles: Implement From and NeuraLayer for NeuraGraph --- src/err.rs | 1 + src/network/graph/from.rs | 61 +++++++++ src/network/graph/mod.rs | 255 +++++++++-------------------------- src/network/graph/node.rs | 57 +++++--- src/network/graph/partial.rs | 196 +++++++++++++++++++++++++++ src/network/residual/axis.rs | 37 +++++ 6 files changed, 404 insertions(+), 203 deletions(-) create mode 100644 src/network/graph/from.rs create mode 100644 src/network/graph/partial.rs diff --git a/src/err.rs b/src/err.rs index 07d9e5c..91a204c 100644 --- a/src/err.rs +++ b/src/err.rs @@ -71,6 +71,7 @@ pub enum NeuraIsolateLayerErr { pub enum NeuraAxisErr { NoInput, ConflictingShape(NeuraShape, NeuraShape), + InvalidAmount(usize, usize, Option), } #[derive(Clone, Debug)] diff --git a/src/network/graph/from.rs b/src/network/graph/from.rs new file mode 100644 index 0000000..166ef57 --- /dev/null +++ b/src/network/graph/from.rs @@ -0,0 +1,61 @@ +use crate::network::residual::{NeuraAxisDefault, NeuraSplitInputs}; + +use super::*; + +trait FromSequential { + fn from_sequential( + seq: &Seq, + nodes: Vec>, + output_shape: NeuraShape, + ) -> Self; +} + +impl FromSequential<(), Data> for NeuraGraph { + fn from_sequential( + _seq: &(), + nodes: Vec>, + output_shape: NeuraShape, + ) -> Self { + Self { + output_index: nodes.len(), + buffer_size: nodes.len() + 1, + nodes: nodes, + output_shape, + } + } +} + +impl< + Data: Clone, + Layer: NeuraLayer + Clone + std::fmt::Debug + 'static, + ChildNetwork, + > FromSequential, Data> for NeuraGraph +where + NeuraGraph: FromSequential, + NeuraAxisDefault: NeuraSplitInputs, +{ + fn from_sequential( + seq: &NeuraSequential, + mut nodes: Vec>, + output_shape: NeuraShape, + ) -> Self { + nodes.push(NeuraGraphNodeConstructed { + node: Box::new(NeuraGraphNode::from(seq.layer.clone())), + inputs: vec![nodes.len()], + output: nodes.len() + 1, + }); + + Self::from_sequential(&seq.child_network, nodes, output_shape) + } +} + +impl From> for NeuraGraph +where + NeuraGraph: FromSequential, Data>, + NeuraSequential: NeuraShapedLayer, +{ + fn from(network: NeuraSequential) -> Self { + let output_shape = network.output_shape(); + Self::from_sequential(&network, vec![], output_shape) + } +} diff --git a/src/network/graph/mod.rs b/src/network/graph/mod.rs index 1c172e1..56dfbdd 100644 --- a/src/network/graph/mod.rs +++ b/src/network/graph/mod.rs @@ -1,131 +1,12 @@ -#![allow(dead_code)] // TODO: remove this - -use std::collections::{HashMap, HashSet, VecDeque}; - -use crate::prelude::*; -use crate::{err::NeuraGraphErr, layer::NeuraShapedLayer}; +use crate::{layer::NeuraShapedLayer, prelude::*}; mod node; pub use node::*; -pub struct NeuraGraphPartial { - pub nodes: Vec>>, - pub output: String, - pub input: String, -} - -#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)] -enum GraphIndex { - Input, - Node(usize), -} - -impl NeuraGraphPartial { - fn get_index_map(&self) -> Result, NeuraGraphErr> { - let mut result = HashMap::with_capacity(self.nodes.len()); - - result.insert(self.input.clone(), GraphIndex::Input); - - for (index, node) in self.nodes.iter().enumerate() { - if result.contains_key(node.name()) { - return Err(NeuraGraphErr::InvalidName(node.name().to_string())); - } - result.insert(node.name().to_string(), GraphIndex::Node(index)); - } - - Ok(result) - } - - fn get_reverse_graph( - &self, - index_map: &HashMap, - ) -> Result>, NeuraGraphErr> { - let mut result = HashMap::new(); - - result.insert(GraphIndex::Input, HashSet::new()); - - for i in 0..self.nodes.len() { - result.insert(GraphIndex::Node(i), HashSet::new()); - } - - for (index, node) in self.nodes.iter().enumerate() { - for input in node.inputs() { - let input_index = index_map - .get(input) - .copied() - .ok_or_else(|| NeuraGraphErr::MissingNode(input.clone()))?; - result - .get_mut(&input_index) - .expect("index_map returned invalid values") - .insert(GraphIndex::Node(index)); - } - } - - Ok(result) - } - - fn get_node_order( - &self, - index_map: &HashMap, - reverse_graph: &HashMap>, - ) -> Result, NeuraGraphErr> { - let mut result: Vec = Vec::new(); - let mut closed: HashSet = HashSet::with_capacity(self.nodes.len()); - let mut open = VecDeque::with_capacity(self.nodes.len()); - open.push_front(GraphIndex::Input); - - /* - index_map.get(&self.output) - .copied() - .ok_or_else(|| NeuraGraphErr::MissingNode(self.output.clone()))? - */ - - while let Some(current) = open.pop_back() { - if closed.contains(¤t) { - continue; - } - - closed.insert(current); - // Do not put 0 (the input) in result - if let GraphIndex::Node(index) = current { - result.push(index); - } - - println!("{:?}", current); - - for next_node in reverse_graph[¤t].iter().copied() { - // Ignore nodes that are already in the closed set - if closed.contains(&next_node) { - continue; - } - - let GraphIndex::Node(node_index) = next_node else { - panic!("Unreachable: cannot have GraphIndex::Input as the output of a node"); - }; - - let inputs = self.nodes[node_index].inputs(); - - // Only consider nodes whose inputs are in the closed set (meaning they would be ready to be evaluated) - if !inputs - .iter() - .all(|input| closed.contains(&index_map[input])) - { - continue; - } - - open.push_front(next_node); - } - } - - if result.len() != self.nodes.len() { - // TODO: verify that if result.len() != self.nodes.len(), then there is a cyclic subgraph - - return Err(NeuraGraphErr::Cyclic); - } +mod partial; +pub use partial::NeuraGraphPartial; - Ok(result) - } -} +mod from; #[derive(Debug)] struct NeuraGraphNodeConstructed { @@ -143,7 +24,7 @@ pub struct NeuraGraph { /// - `nodes[0].inputs = [0]` nodes: Vec>, - input_shape: NeuraShape, + // input_shape: NeuraShape, output_shape: NeuraShape, output_index: usize, @@ -156,82 +37,57 @@ impl NeuraShapedLayer for NeuraGraph { } } -impl NeuraPartialLayer for NeuraGraphPartial { - type Constructed = NeuraGraph; +impl NeuraGraph { + fn create_buffer(&self) -> Vec> { + let mut res = Vec::with_capacity(self.buffer_size); - type Err = NeuraGraphErr; + for _ in 0..self.buffer_size { + res.push(None); + } + + res + } - fn construct(self, input_shape: NeuraShape) -> Result { - let index_map = self.get_index_map()?; - let reverse_graph = self.get_reverse_graph(&index_map)?; + fn eval_in(&self, input: &Data, buffer: &mut Vec>) + where + Data: Clone, + { + buffer[0] = Some(input.clone()); - // List out the nodes in their execution order - let node_order = self.get_node_order(&index_map, &reverse_graph)?; - let mut new_index_map: HashMap = HashMap::from_iter( - node_order + for node in self.nodes.iter() { + // PERF: re-use the allocation for `inputs`, and `.take()` the elements only needed once? + let inputs: Vec<_> = node + .inputs .iter() - .map(|&i| (self.nodes[i].name().to_string(), i)), - ); - new_index_map.insert(self.input.clone(), 0); - - // TODO: filter out the nodes that are not necessary for computing the result (BFS from the output node back to the inputs) - // A temporary solution can be to trim the graph - let output_index = new_index_map - .get(&self.output) - .copied() - .ok_or_else(|| NeuraGraphErr::MissingNode(self.output.clone()))?; - - let mut nodes = Vec::with_capacity(self.nodes.len()); - let mut shapes: Vec> = vec![None; self.nodes.len() + 1]; - shapes[0] = Some(input_shape); - - for index in node_order.into_iter() { - let node = &*self.nodes[index]; - let node_inputs = node.inputs(); - let mut inputs = Vec::with_capacity(node_inputs.len()); - let mut input_shapes = Vec::with_capacity(node_inputs.len()); - - for input in node_inputs { - let input_index = new_index_map.get(input).copied().expect( - "Unreachable: new_index_map should contain all nodes defined and all nodes should have existing nodes as input" - ); - inputs.push(input_index); - input_shapes.push(shapes[input_index].expect( - "Unreachable: the order of execution should guarantee that all inputs have appeared before") - ); - } - - let (constructed, output_shape) = node - .construct(input_shapes) - .map_err(|e| NeuraGraphErr::LayerErr(e))?; - - shapes[index] = Some(output_shape); - - nodes.push(NeuraGraphNodeConstructed { - node: constructed, - inputs, - output: new_index_map - .get(node.name()) - .copied() - .unwrap_or_else(|| unreachable!()), - }); + .map(|&i| { + buffer[i] + .clone() + .expect("Unreachable: output of previous layer was not set") + }) + .collect(); + let result = node.node.eval(&inputs); + buffer[node.output] = Some(result); } + } +} + +impl NeuraLayer for NeuraGraph { + type Output = Data; + + fn eval(&self, input: &Data) -> Self::Output { + let mut buffer = self.create_buffer(); - let output_shape = shapes[output_index].unwrap_or_else(|| unreachable!()); + self.eval_in(input, &mut buffer); - Ok(NeuraGraph { - nodes, - input_shape, - output_shape, - output_index, - buffer_size: self.nodes.len() + 1, - }) + buffer[self.output_index] + .take() + .expect("Unreachable: output was not set") } } #[cfg(test)] mod test { - use crate::network::residual::NeuraAxisAppend; + use crate::{err::NeuraGraphErr, network::residual::NeuraAxisAppend, utils::uniform_vector}; use super::*; @@ -346,4 +202,27 @@ mod test { NeuraGraphErr::MissingNode(String::from("missing")) ); } + + #[test] + fn test_eval_equal_sequential() { + let network = neura_sequential![ + neura_layer!("dense", 4, f64), + neura_layer!("dense", 2, f64), + neura_layer!("softmax") + ] + .construct(NeuraShape::Vector(3)) + .unwrap(); + + let graph = NeuraGraph::from(network.clone()); + + for _ in 0..10 { + let input = uniform_vector(3); + let seq_result = network.eval(&input); + let graph_result = graph.eval(&input); + + assert_eq!(seq_result.shape(), graph_result.shape()); + approx::assert_relative_eq!(seq_result[0], graph_result[0]); + approx::assert_relative_eq!(seq_result[1], graph_result[1]); + } + } } diff --git a/src/network/graph/node.rs b/src/network/graph/node.rs index afd8d4b..0635739 100644 --- a/src/network/graph/node.rs +++ b/src/network/graph/node.rs @@ -1,14 +1,15 @@ use dyn_clone::DynClone; +use std::fmt::Debug; use crate::{ err::NeuraAxisErr, layer::{NeuraLayer, NeuraShapedLayer}, - network::residual::{NeuraCombineInputs, NeuraSplitInputs}, + network::residual::{NeuraAxisDefault, NeuraCombineInputs, NeuraSplitInputs}, prelude::{NeuraPartialLayer, NeuraShape}, }; // TODO: split into two traits -pub trait NeuraGraphNodePartial: DynClone + std::fmt::Debug { +pub trait NeuraGraphNodePartial: DynClone + Debug { fn inputs<'a>(&'a self) -> &'a [String]; fn name<'a>(&'a self) -> &'a str; @@ -18,7 +19,7 @@ pub trait NeuraGraphNodePartial: DynClone + std::fmt::Debug { ) -> Result<(Box>, NeuraShape), String>; } -pub trait NeuraGraphNodeEval: DynClone + std::fmt::Debug { +pub trait NeuraGraphNodeEval: DynClone + Debug { fn eval<'a>(&'a self, inputs: &[Data]) -> Data; } @@ -46,15 +47,15 @@ impl NeuraGraphNode { Axis: NeuraSplitInputs + NeuraCombineInputs> + Clone - + std::fmt::Debug + + Debug + 'static, - Layer: NeuraPartialLayer + Clone + std::fmt::Debug + 'static, + Layer: NeuraPartialLayer + Clone + Debug + 'static, Layer::Constructed: NeuraShapedLayer + NeuraLayer<>::Combined, Output = Data> + Clone - + std::fmt::Debug + + Debug + 'static, - Layer::Err: std::fmt::Debug, + Layer::Err: Debug, { Box::new(self) } @@ -62,10 +63,8 @@ impl NeuraGraphNode { impl< Data: Clone, - Axis: NeuraSplitInputs + Clone + std::fmt::Debug, - Layer: NeuraLayer<>::Combined, Output = Data> - + Clone - + std::fmt::Debug, + Axis: NeuraSplitInputs + Clone + Debug, + Layer: NeuraLayer<>::Combined, Output = Data> + Clone + Debug, > NeuraGraphNodeEval for NeuraGraphNode { fn eval<'a>(&'a self, inputs: &[Data]) -> Data { @@ -75,22 +74,33 @@ impl< } } +impl From for NeuraGraphNode { + fn from(layer: Layer) -> Self { + Self { + inputs: vec![], + axis: NeuraAxisDefault, + layer, + name: random_name(), + } + } +} + impl< Data: Clone, Axis: NeuraSplitInputs + NeuraCombineInputs> + Clone - + std::fmt::Debug + + Debug + 'static, - Layer: NeuraPartialLayer + Clone + std::fmt::Debug, + Layer: NeuraPartialLayer + Clone + Debug, > NeuraGraphNodePartial for NeuraGraphNode where Layer::Constructed: NeuraShapedLayer + NeuraLayer<>::Combined, Output = Data> + Clone - + std::fmt::Debug + + Debug + 'static, - Layer::Err: std::fmt::Debug, + Layer::Err: Debug, { fn inputs<'a>(&'a self) -> &'a [String] { &self.inputs @@ -127,3 +137,20 @@ where )) } } + +pub fn random_name() -> String { + use rand::Rng; + use std::fmt::Write; + + let mut res = String::with_capacity(10); + write!(&mut res, "value_").unwrap(); + + let mut rng = rand::thread_rng(); + + for _ in 0..4 { + let ch = char::from_u32(rng.gen_range((b'a' as u32)..(b'z' as u32))).unwrap(); + write!(&mut res, "{}", ch).unwrap(); + } + + res +} diff --git a/src/network/graph/partial.rs b/src/network/graph/partial.rs new file mode 100644 index 0000000..99ad1d8 --- /dev/null +++ b/src/network/graph/partial.rs @@ -0,0 +1,196 @@ +use crate::err::NeuraGraphErr; +use std::collections::{HashMap, HashSet, VecDeque}; + +use super::*; + +pub struct NeuraGraphPartial { + pub nodes: Vec>>, + pub output: String, + pub input: String, +} + +#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)] +pub(crate) enum GraphIndex { + Input, + Node(usize), +} + +impl NeuraGraphPartial { + pub(crate) fn get_index_map(&self) -> Result, NeuraGraphErr> { + let mut result = HashMap::with_capacity(self.nodes.len()); + + result.insert(self.input.clone(), GraphIndex::Input); + + for (index, node) in self.nodes.iter().enumerate() { + if result.contains_key(node.name()) { + return Err(NeuraGraphErr::InvalidName(node.name().to_string())); + } + result.insert(node.name().to_string(), GraphIndex::Node(index)); + } + + Ok(result) + } + + pub(crate) fn get_reverse_graph( + &self, + index_map: &HashMap, + ) -> Result>, NeuraGraphErr> { + let mut result = HashMap::new(); + + result.insert(GraphIndex::Input, HashSet::new()); + + for i in 0..self.nodes.len() { + result.insert(GraphIndex::Node(i), HashSet::new()); + } + + for (index, node) in self.nodes.iter().enumerate() { + for input in node.inputs() { + let input_index = index_map + .get(input) + .copied() + .ok_or_else(|| NeuraGraphErr::MissingNode(input.clone()))?; + result + .get_mut(&input_index) + .expect("index_map returned invalid values") + .insert(GraphIndex::Node(index)); + } + } + + Ok(result) + } + + pub(crate) fn get_node_order( + &self, + index_map: &HashMap, + reverse_graph: &HashMap>, + ) -> Result, NeuraGraphErr> { + let mut result: Vec = Vec::new(); + let mut closed: HashSet = HashSet::with_capacity(self.nodes.len()); + let mut open = VecDeque::with_capacity(self.nodes.len()); + open.push_front(GraphIndex::Input); + + /* + index_map.get(&self.output) + .copied() + .ok_or_else(|| NeuraGraphErr::MissingNode(self.output.clone()))? + */ + + while let Some(current) = open.pop_back() { + if closed.contains(¤t) { + continue; + } + + closed.insert(current); + // Do not put 0 (the input) in result + if let GraphIndex::Node(index) = current { + result.push(index); + } + + println!("{:?}", current); + + for next_node in reverse_graph[¤t].iter().copied() { + // Ignore nodes that are already in the closed set + if closed.contains(&next_node) { + continue; + } + + let GraphIndex::Node(node_index) = next_node else { + panic!("Unreachable: cannot have GraphIndex::Input as the output of a node"); + }; + + let inputs = self.nodes[node_index].inputs(); + + // Only consider nodes whose inputs are in the closed set (meaning they would be ready to be evaluated) + if !inputs + .iter() + .all(|input| closed.contains(&index_map[input])) + { + continue; + } + + open.push_front(next_node); + } + } + + if result.len() != self.nodes.len() { + // TODO: verify that if result.len() != self.nodes.len(), then there is a cyclic subgraph + + return Err(NeuraGraphErr::Cyclic); + } + + Ok(result) + } +} + +impl NeuraPartialLayer for NeuraGraphPartial { + type Constructed = NeuraGraph; + + type Err = NeuraGraphErr; + + fn construct(self, input_shape: NeuraShape) -> Result { + let index_map = self.get_index_map()?; + let reverse_graph = self.get_reverse_graph(&index_map)?; + + // List out the nodes in their execution order + let node_order = self.get_node_order(&index_map, &reverse_graph)?; + let mut new_index_map: HashMap = HashMap::from_iter( + node_order + .iter() + .map(|&i| (self.nodes[i].name().to_string(), i)), + ); + new_index_map.insert(self.input.clone(), 0); + + // TODO: filter out the nodes that are not necessary for computing the result (BFS from the output node back to the inputs) + // A temporary solution can be to trim the graph + let output_index = new_index_map + .get(&self.output) + .copied() + .ok_or_else(|| NeuraGraphErr::MissingNode(self.output.clone()))?; + + let mut nodes = Vec::with_capacity(self.nodes.len()); + let mut shapes: Vec> = vec![None; self.nodes.len() + 1]; + shapes[0] = Some(input_shape); + + for index in node_order.into_iter() { + let node = &*self.nodes[index]; + let node_inputs = node.inputs(); + let mut inputs = Vec::with_capacity(node_inputs.len()); + let mut input_shapes = Vec::with_capacity(node_inputs.len()); + + for input in node_inputs { + let input_index = new_index_map.get(input).copied().expect( + "Unreachable: new_index_map should contain all nodes defined and all nodes should have existing nodes as input" + ); + inputs.push(input_index); + input_shapes.push(shapes[input_index].expect( + "Unreachable: the order of execution should guarantee that all inputs have appeared before") + ); + } + + let (constructed, output_shape) = node + .construct(input_shapes) + .map_err(|e| NeuraGraphErr::LayerErr(e))?; + + shapes[index] = Some(output_shape); + + nodes.push(NeuraGraphNodeConstructed { + node: constructed, + inputs, + output: new_index_map + .get(node.name()) + .copied() + .unwrap_or_else(|| unreachable!()), + }); + } + + let output_shape = shapes[output_index].unwrap_or_else(|| unreachable!()); + + Ok(NeuraGraph { + nodes, + // input_shape, + output_shape, + output_index, + buffer_size: self.nodes.len() + 1, + }) + } +} diff --git a/src/network/residual/axis.rs b/src/network/residual/axis.rs index 9baa3cf..9e34c8d 100644 --- a/src/network/residual/axis.rs +++ b/src/network/residual/axis.rs @@ -4,6 +4,8 @@ use nalgebra::{Const, DVector, Dyn, Scalar, VecStorage}; use crate::{err::NeuraAxisErr, prelude::NeuraShape}; +// TODO: create a NeuraAxis trait + #[derive(Clone, Copy, Debug)] pub struct NeuraAxisAppend; @@ -34,6 +36,7 @@ impl NeuraCombineInputs> for NeuraAxisAppend { } } +// TODO: use another trait for combining NeuraShape, or make it another member of the trait impl NeuraCombineInputs for NeuraAxisAppend { type Combined = Result; @@ -80,3 +83,37 @@ impl NeuraSplitInputs> for NeuraAxisAppe result } } + +#[derive(Clone, Debug)] +pub struct NeuraAxisDefault; + +impl NeuraCombineInputs> for NeuraAxisDefault { + type Combined = DVector; + + fn combine(&self, inputs: Vec>>) -> Self::Combined { + assert!(inputs.len() == 1); + + inputs[0].borrow().clone() + } +} + +impl NeuraCombineInputs for NeuraAxisDefault { + type Combined = Result; + + fn combine(&self, inputs: Vec>) -> Self::Combined { + if inputs.len() != 1 { + Err(NeuraAxisErr::InvalidAmount(inputs.len(), 1, Some(1))) + } else { + Ok(*inputs[0].borrow()) + } + } +} + +impl NeuraSplitInputs for NeuraAxisDefault +where + NeuraAxisDefault: NeuraCombineInputs, +{ + fn split(&self, combined: &Self::Combined, _input_shapes: &[NeuraShape]) -> Vec { + vec![combined.clone()] + } +}