microtensor 0.1.5

Automatic differentiation for tensor operations
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use std::io;
use std::fs;
use std::collections::HashMap;

use itertools::Itertools;
use serde::{Serialize, Deserialize, de::DeserializeOwned};

use crate::{
  internal::*,
  scalar::Real,
  variable::{ Variable, Node, NodeCell, Op }
};


/// Snapshot of a computation graph with defined inputs and outputs.
///
/// Can be used for recomputing the entire graph or saving it to disc.

#[derive(Debug, Clone)]
pub struct Graph<T: Real + 'static> {
  pub inputs: Vec<Variable<T>>,
  pub outputs: Vec<Variable<T>>,
}

impl<T: Real + Serialize + DeserializeOwned + 'static> Graph<T> {
  pub fn new(inputs: &[Variable<T>], outputs: &[Variable<T>]) -> Self {
    Self {
      inputs: inputs.into(),
      outputs: outputs.into(),
    }
  }

  pub fn run(&self, inputs: &[&Variable<T>]) {
    for (input, data) in self.inputs.iter().zip(inputs) {
      input.feed(data);
    }
    for output in &self.outputs {
      output.forward(); //XXX Don't calculate entire graph multiple times
    }
  }

  fn history(&self) -> Vec<RcT<Node<T>>> {
    let mut history = self.outputs
      .iter()
      .map(|out| out.history() )
      .collect::<Vec<_>>()
      .concat();
    history.sort_by(|a, b| a.id.partial_cmp(&b.id).unwrap() );
    history.into_iter().unique_by(|a| a.id ).collect()
  }

  pub fn parameters(&self) -> Vec<Variable<T>> {
    self.history()
      .into_iter()
      .filter(|node| node.trainable )
      .map(|node| Variable { node } )
      .collect()
  }

  pub fn load(filename: &str) -> io::Result<Self> {
    let bytes = fs::read(filename)?;
    let tensor_dump: GraphDump<T> = postcard::from_bytes(&bytes).unwrap();
    let mut nodes: HashMap<usize, RcT<Node<T>>> = HashMap::new();
    for dump in tensor_dump.history {
      let node = Node {
        id: dump.id,
        cell: dump.cell.clone(),
        op: dump.op,
        previous: dump.previous.iter().map(|id| nodes[id].clone() ).collect(),
        trainable: dump.trainable,
      };
      nodes.insert(node.id, RcT::new(node));
    }
    let map_tensor = |dump: VariableDump| Variable {
      node: nodes[&dump.node].clone(),
    };
    Ok(Graph {
      inputs: tensor_dump.inputs.into_iter().map(map_tensor).collect(),
      outputs: tensor_dump.outputs.into_iter().map(map_tensor).collect(),
    })
  }

  pub fn save(&self, filename: &str) -> io::Result<()> {
    let history = self.history();

    let history_dump = history.iter().map(|node| {
      let op: Vec<u8> = postcard::to_allocvec(&node.op).unwrap();
      let op = postcard::from_bytes(&op).unwrap();
      let mut cell = node.cell.clone();
      cell.data = cell.data.detach();
      cell.grad = cell.grad.and_then(|grad| Some(grad.detach()) );
      NodeDump {
        id: node.id,
        cell,
        op,
        previous: node.previous.iter().map(|prev| prev.id ).collect(),
        trainable: node.trainable,
      }
    }).collect();

    let map_tensor = |tensor: &Variable<T>| VariableDump {
      node: tensor.node.id,
    };

    let graph_dump = GraphDump {
      history: history_dump,
      inputs: self.inputs.iter().map(map_tensor).collect(),
      outputs: self.outputs.iter().map(map_tensor).collect(),
    };

    let data: Vec<u8> = postcard::to_allocvec(&graph_dump).unwrap();
    fs::write(filename, data)
  }
}


#[derive(Serialize, Deserialize)]
struct NodeDump<T: Real + 'static> {
  id: usize,
  cell: NodeCell<T>, //XXX don't save gradients
  op: Option<Op<T>>,
  previous: Vec<usize>,
  trainable: bool,
}

#[derive(Serialize, Deserialize)]
struct VariableDump {
  node: usize,
}

#[derive(Serialize, Deserialize)]
struct GraphDump<T: Real + 'static> {
  history: Vec<NodeDump<T>>,
  inputs: Vec<VariableDump>,
  outputs: Vec<VariableDump>,
}