tfmicro 0.1.0

Rust bindings of TensorFlow Lite for Microcontrollers
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

//! Rust Bindings for Tensor type

use core::convert::{TryFrom, TryInto};
use core::mem::size_of;
use core::slice;
use ordered_float::NotNan;

use crate::bindings;

mod info;
pub use info::TensorInfo;

/// A TensorFlow Tensor
#[repr(transparent)]
#[derive(Default)]
pub struct Tensor(bindings::TfLiteTensor);

/// Tensor element data types supported
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum ElementType {
    Float32,
    UInt8,
    Int32,
}
impl TryFrom<bindings::TfLiteType> for ElementType {
    type Error = bindings::TfLiteType;

    fn try_from(status: bindings::TfLiteType) -> Result<Self, Self::Error> {
        use ElementType::*;

        match status {
            bindings::TfLiteType::kTfLiteFloat32 => Ok(Float32),
            bindings::TfLiteType::kTfLiteUInt8 => Ok(UInt8),
            bindings::TfLiteType::kTfLiteInt32 => Ok(Int32),
            t => Err(t),
        }
    }
}

/// Marker trait for those intristic types we support
pub trait ElemTypeOf {
    fn elem_type_of() -> ElementType;
}

impl ElemTypeOf for NotNan<f32> {
    fn elem_type_of() -> ElementType {
        ElementType::Float32
    }
}
impl ElemTypeOf for f32 {
    fn elem_type_of() -> ElementType {
        ElementType::Float32
    }
}
impl ElemTypeOf for u8 {
    fn elem_type_of() -> ElementType {
        ElementType::UInt8
    }
}
impl ElemTypeOf for i32 {
    fn elem_type_of() -> ElementType {
        ElementType::Int32
    }
}

/// Implement From raw types to Tensor
impl From<*mut bindings::TfLiteTensor> for &Tensor {
    fn from(t: *mut bindings::TfLiteTensor) -> Self {
        unsafe { &*(t as *mut Tensor) }
    }
}
impl From<*mut bindings::TfLiteTensor> for &mut Tensor {
    fn from(t: *mut bindings::TfLiteTensor) -> Self {
        unsafe { &mut *(t as *mut Tensor) }
    }
}

impl Tensor {
    /// The element type of this tensor.
    ///
    /// Returns `Some(element_type)` if the element type annotated on this
    /// tensor matches a member of
    /// [`ElementType`](crate::tensor::ElementType). Otherwise returns `None`.
    pub fn element_type(&self) -> Option<ElementType> {
        self.0.type_.try_into().ok()
    }

    /// Returns a [`TensorInfo`](crate::tensor::TensorInfo) that describes this tensor
    ///
    /// # Panics
    ///
    /// Panics if the underlying tensor cannot be represented by a
    /// [`TensorInfo`](crate::tensor::TensorInfo).
    pub fn info(&self) -> TensorInfo {
        self.inner().try_into().unwrap()
    }

    pub(crate) fn inner(&self) -> &bindings::TfLiteTensor {
        &self.0
    }

    /// Extracts the tensor's data as a flat slice.
    ///
    /// Call the [info](#method.info) method to check the dimensionality of
    /// the tensor.
    ///
    /// # Panics
    ///
    /// This method will panic if `T` does not match the data type
    /// annotated on this tensor. Call
    /// [element_type()](#method.element_type) to discover the data type.
    pub fn as_data<T>(&self) -> &[T]
    where
        T: ElemTypeOf,
    {
        assert!(
            self.element_type().unwrap() == T::elem_type_of(),
            "Type `{:?}` does not match the original type `{:?}`",
            T::elem_type_of(),
            self.0.type_
        );

        unsafe {
            slice::from_raw_parts(
                self.0.data.raw_const as *const T,
                self.0.bytes / size_of::<T>(),
            )
        }
    }

    /// Extracts the tensor's data as a mutable flat slice.
    ///
    /// Call the [info](#method.info) method to check the dimensionality of
    /// the tensor.
    ///
    /// # Panics
    ///
    /// This method will panic if `T` does not match the data type
    /// annotated on this tensor. Call
    /// [element_type()](#method.element_type) to discover the data type.
    pub fn as_data_mut<T>(&mut self) -> &mut [T]
    where
        T: ElemTypeOf,
    {
        assert!(
            self.element_type().unwrap() == T::elem_type_of(),
            "Type `{:?}` does not match the original type `{:?}`",
            T::elem_type_of(),
            self.0.type_
        );

        unsafe {
            slice::from_raw_parts_mut(
                self.0.data.raw as *mut T,
                self.0.bytes / size_of::<T>(),
            )
        }
    }
}