parent
efaed91f83
commit
251e4d02d2
@ -0,0 +1,61 @@
|
||||
use crate::network::residual::{NeuraAxisDefault, NeuraSplitInputs};
|
||||
|
||||
use super::*;
|
||||
|
||||
trait FromSequential<Seq, Data> {
|
||||
fn from_sequential(
|
||||
seq: &Seq,
|
||||
nodes: Vec<NeuraGraphNodeConstructed<Data>>,
|
||||
output_shape: NeuraShape,
|
||||
) -> Self;
|
||||
}
|
||||
|
||||
impl<Data> FromSequential<(), Data> for NeuraGraph<Data> {
|
||||
fn from_sequential(
|
||||
_seq: &(),
|
||||
nodes: Vec<NeuraGraphNodeConstructed<Data>>,
|
||||
output_shape: NeuraShape,
|
||||
) -> Self {
|
||||
Self {
|
||||
output_index: nodes.len(),
|
||||
buffer_size: nodes.len() + 1,
|
||||
nodes: nodes,
|
||||
output_shape,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<
|
||||
Data: Clone,
|
||||
Layer: NeuraLayer<Data, Output = Data> + Clone + std::fmt::Debug + 'static,
|
||||
ChildNetwork,
|
||||
> FromSequential<NeuraSequential<Layer, ChildNetwork>, Data> for NeuraGraph<Data>
|
||||
where
|
||||
NeuraGraph<Data>: FromSequential<ChildNetwork, Data>,
|
||||
NeuraAxisDefault: NeuraSplitInputs<Data, Combined = Data>,
|
||||
{
|
||||
fn from_sequential(
|
||||
seq: &NeuraSequential<Layer, ChildNetwork>,
|
||||
mut nodes: Vec<NeuraGraphNodeConstructed<Data>>,
|
||||
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<Data, Layer, ChildNetwork> From<NeuraSequential<Layer, ChildNetwork>> for NeuraGraph<Data>
|
||||
where
|
||||
NeuraGraph<Data>: FromSequential<NeuraSequential<Layer, ChildNetwork>, Data>,
|
||||
NeuraSequential<Layer, ChildNetwork>: NeuraShapedLayer,
|
||||
{
|
||||
fn from(network: NeuraSequential<Layer, ChildNetwork>) -> Self {
|
||||
let output_shape = network.output_shape();
|
||||
Self::from_sequential(&network, vec![], output_shape)
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,196 @@
|
||||
use crate::err::NeuraGraphErr;
|
||||
use std::collections::{HashMap, HashSet, VecDeque};
|
||||
|
||||
use super::*;
|
||||
|
||||
pub struct NeuraGraphPartial<Data> {
|
||||
pub nodes: Vec<Box<dyn NeuraGraphNodePartial<Data>>>,
|
||||
pub output: String,
|
||||
pub input: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash)]
|
||||
pub(crate) enum GraphIndex {
|
||||
Input,
|
||||
Node(usize),
|
||||
}
|
||||
|
||||
impl<Data> NeuraGraphPartial<Data> {
|
||||
pub(crate) fn get_index_map(&self) -> Result<HashMap<String, GraphIndex>, 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<String, GraphIndex>,
|
||||
) -> Result<HashMap<GraphIndex, HashSet<GraphIndex>>, 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<String, GraphIndex>,
|
||||
reverse_graph: &HashMap<GraphIndex, HashSet<GraphIndex>>,
|
||||
) -> Result<Vec<usize>, NeuraGraphErr> {
|
||||
let mut result: Vec<usize> = Vec::new();
|
||||
let mut closed: HashSet<GraphIndex> = 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<Data> NeuraPartialLayer for NeuraGraphPartial<Data> {
|
||||
type Constructed = NeuraGraph<Data>;
|
||||
|
||||
type Err = NeuraGraphErr;
|
||||
|
||||
fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
|
||||
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<String, usize> = 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<Option<NeuraShape>> = 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,
|
||||
})
|
||||
}
|
||||
}
|
||||
Loading…
Reference in new issue