tract-tensorflow 0.2.1

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
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

use std::collections::HashMap;

use tract_core::ops::Op;
use tract_core::TractResult;

use crate::tfpb::node_def::NodeDef;

#[macro_use]
mod macros;

pub mod array;
pub mod logic;
pub mod math;
pub mod nn;
pub mod quant;

pub type OpRegister = HashMap<&'static str, fn(&NodeDef) -> TractResult<Box<Op>>>;

pub struct OpBuilder(OpRegister);

impl OpBuilder {
    pub fn new() -> OpBuilder {
        let mut reg = OpRegister::new();
        array::register_all_ops(&mut reg);
        logic::register_all_ops(&mut reg);
        math::register_all_ops(&mut reg);
        nn::register_all_ops(&mut reg);
        quant::register_all_ops(&mut reg);
        reg.insert("Const", konst);
        reg.insert("Placeholder", placeholder);
        OpBuilder(reg)
    }

    pub fn build(&self, pb: &NodeDef) -> TractResult<Box<Op>> {
        match self.0.get(pb.get_op()) {
            Some(builder) => builder(pb),
            None => Ok(Box::new(::tract_core::ops::unimpl::UnimplementedOp::new(
                pb.get_op(),
                format!("{:?}", pb),
            ))),
        }
    }
}

pub fn konst(node: &NodeDef) -> TractResult<Box<Op>> {
    let dtype = node.get_attr_datum_type("dtype")?;
    let mat = node.get_attr_tensor("value")?;

    if mat.datum_type() != dtype {
        bail!(
            "Const node {:?} doesn't have the expected {:?} type.",
            mat,
            dtype
        );
    }

    Ok(Box::new(::tract_core::ops::konst::Const::for_tensor(mat)))
}

pub fn placeholder(node: &NodeDef) -> TractResult<Box<Op>> {
    let dt = node.get_attr_datum_type("dtype")?;
    Ok(Box::new(::tract_core::ops::source::Source::new(
        ::tract_core::analyser::types::TensorFact::dt(dt.into()),
    )))
}