tract_tensorflow/
tensor.rs

1use tract_hir::internal::*;
2
3use crate::tfpb::tensorflow::tensor_shape_proto::Dim;
4use crate::tfpb::tensorflow::{TensorProto, TensorShapeProto};
5
6use crate::tfpb::tensorflow::DataType;
7use std::convert::TryFrom;
8
9impl TryFrom<DataType> for DatumType {
10    type Error = TractError;
11    fn try_from(t: DataType) -> TractResult<DatumType> {
12        match t {
13            DataType::DtBool => Ok(DatumType::Bool),
14            DataType::DtUint8 => Ok(DatumType::U8),
15            DataType::DtUint16 => Ok(DatumType::U16),
16            DataType::DtUint32 => Ok(DatumType::U32),
17            DataType::DtUint64 => Ok(DatumType::U64),
18            DataType::DtInt8 => Ok(DatumType::I8),
19            DataType::DtInt16 => Ok(DatumType::I16),
20            DataType::DtInt32 => Ok(DatumType::I32),
21            DataType::DtInt64 => Ok(DatumType::I64),
22            DataType::DtHalf => Ok(DatumType::F16),
23            DataType::DtFloat => Ok(DatumType::F32),
24            DataType::DtDouble => Ok(DatumType::F64),
25            DataType::DtString => Ok(DatumType::Blob),
26            _ => Err(format_err!("Unknown DatumType {:?}", t))?,
27        }
28    }
29}
30
31impl<'a> TryFrom<&'a TensorShapeProto> for TVec<isize> {
32    type Error = TractError;
33    fn try_from(t: &'a TensorShapeProto) -> TractResult<TVec<isize>> {
34        Ok(t.dim.iter().map(|d| d.size as isize).collect::<TVec<_>>())
35    }
36}
37
38impl<'a> TryFrom<&'a TensorShapeProto> for TVec<usize> {
39    type Error = TractError;
40    fn try_from(t: &'a TensorShapeProto) -> TractResult<TVec<usize>> {
41        if t.dim.iter().any(|d| d.size < 0) {
42            bail!("Negative dim found")
43        }
44        Ok(t.dim.iter().map(|d| d.size as usize).collect::<TVec<_>>())
45    }
46}
47
48impl TryFrom<DatumType> for DataType {
49    type Error = TractError;
50    fn try_from(dt: DatumType) -> TractResult<DataType> {
51        match dt {
52            DatumType::Bool => Ok(DataType::DtBool),
53            DatumType::U8 => Ok(DataType::DtUint8),
54            DatumType::U16 => Ok(DataType::DtUint16),
55            DatumType::U32 => Ok(DataType::DtUint32),
56            DatumType::U64 => Ok(DataType::DtUint64),
57            DatumType::I8 => Ok(DataType::DtInt8),
58            DatumType::I16 => Ok(DataType::DtInt16),
59            DatumType::I32 => Ok(DataType::DtInt32),
60            DatumType::I64 => Ok(DataType::DtInt64),
61            DatumType::F16 => Ok(DataType::DtHalf),
62            DatumType::F32 => Ok(DataType::DtFloat),
63            DatumType::F64 => Ok(DataType::DtDouble),
64            DatumType::Blob => Ok(DataType::DtString),
65            DatumType::String => Ok(DataType::DtString),
66            DatumType::QI8(_) => Ok(DataType::DtQint8),
67            DatumType::QU8(_) => Ok(DataType::DtQint8),
68            DatumType::QI32(_) => Ok(DataType::DtQint32),
69            _ => bail!("DatumType is not translatable in protobuf"),
70        }
71    }
72}
73
74fn tensor_from_repeated_field<T: Datum>(shape: &[usize], data: Vec<T>) -> TractResult<Tensor> {
75    let t = if data.len() == 1 {
76        tract_ndarray::ArrayD::from_elem(shape, data[0].clone()).into()
77    } else {
78        tract_ndarray::ArrayD::from_shape_vec(shape, data.to_vec())?.into()
79    };
80    Ok(t)
81}
82
83impl TryFrom<&TensorProto> for Tensor {
84    type Error = TractError;
85    fn try_from(t: &TensorProto) -> TractResult<Tensor> {
86        let shape = t
87            .tensor_shape
88            .as_ref()
89            .ok_or_else(|| format_err!("TensorProto is missing its tensor_shape"))?;
90        // Go through the checked conversion instead of `as _`: it rejects
91        // negative dims instead of wrapping them into huge `usize` values.
92        let dims: TVec<usize> = shape.try_into()?;
93        let rank = dims.len();
94        let content = &t.tensor_content;
95        let dtype = DataType::try_from(t.dtype)?;
96        let mat: Tensor = if content.len() != 0 {
97            unsafe {
98                match dtype {
99                    DataType::DtFloat => Self::from_raw::<f32>(&dims, content)?,
100                    DataType::DtDouble => Self::from_raw::<f64>(&dims, content)?,
101                    DataType::DtInt32 => Self::from_raw::<i32>(&dims, content)?,
102                    DataType::DtInt64 => Self::from_raw::<i64>(&dims, content)?,
103                    _ => bail!("unsupported tf DataType: {:?}", dtype),
104                }
105            }
106        } else {
107            match dtype {
108                DataType::DtInt32 => tensor_from_repeated_field(&dims, t.int_val.to_vec())?,
109                DataType::DtInt64 => tensor_from_repeated_field(&dims, t.int64_val.to_vec())?,
110                DataType::DtFloat => tensor_from_repeated_field(&dims, t.float_val.to_vec())?,
111                DataType::DtDouble => tensor_from_repeated_field(&dims, t.double_val.to_vec())?,
112                DataType::DtString => {
113                    let strings = t
114                        .string_val
115                        .iter()
116                        .map(|s| Blob::try_from(&**s))
117                        .collect::<TractResult<Vec<Blob>>>()?;
118                    tensor_from_repeated_field(&dims, strings)?
119                }
120                _ => bail!("unsupported tf DataType: {:?}", t.dtype),
121            }
122        };
123        assert_eq!(rank, mat.shape().len());
124        Ok(mat)
125    }
126}
127
128fn empty_tensor_proto() -> TensorProto {
129    TensorProto {
130        dtype: 0,
131        tensor_shape: None,
132        version_number: 0,
133        tensor_content: vec![],
134        half_val: vec![],
135        float_val: vec![],
136        double_val: vec![],
137        int_val: vec![],
138        string_val: vec![],
139        scomplex_val: vec![],
140        dcomplex_val: vec![],
141        resource_handle_val: vec![],
142        variant_val: vec![],
143        uint32_val: vec![],
144        uint64_val: vec![],
145        int64_val: vec![],
146        bool_val: vec![],
147    }
148}
149
150impl TryFrom<&Tensor> for TensorProto {
151    type Error = TractError;
152    fn try_from(from: &Tensor) -> TractResult<TensorProto> {
153        let mut tensor = empty_tensor_proto();
154        let shape = TensorShapeProto {
155            dim: from.shape().iter().map(|d| Dim { size: *d as _, name: String::new() }).collect(),
156            unknown_rank: false,
157        };
158        tensor.tensor_shape = Some(shape);
159        let dt = DataType::try_from(from.datum_type())?;
160        tensor.dtype = dt.into();
161        match from.datum_type() {
162            DatumType::F32 => {
163                tensor.float_val = from.to_plain_array_view::<f32>()?.iter().cloned().collect();
164            }
165            DatumType::F64 => {
166                tensor.double_val = from.to_plain_array_view::<f64>()?.iter().cloned().collect();
167            }
168            DatumType::I32 => {
169                tensor.int_val = from.to_plain_array_view::<i32>()?.iter().cloned().collect();
170            }
171            DatumType::I64 => {
172                tensor.int64_val = from.to_plain_array_view::<i64>()?.iter().cloned().collect();
173            }
174            _ => bail!("unsupported datum type for protobuf export: {:?}", from.datum_type()),
175        }
176        Ok(tensor)
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    // Regression: malformed `TensorProto` fields used to abort the process.
185    // `tensor_shape` and `dtype` were `unwrap()`-ed, and dtypes tract does not
186    // support fell into `unimplemented!()`. The ONNX loader was fixed the same
187    // way (see `onnx/src/tensor.rs`).
188    #[test]
189    fn malformed_tensor_proto_does_not_panic() {
190        // dtype that prost does not recognise
191        let unknown_dtype = TensorProto { dtype: 9999, ..empty_tensor_proto() };
192        assert!(Tensor::try_from(&unknown_dtype).is_err());
193
194        // no tensor_shape at all
195        let no_shape = TensorProto { dtype: DataType::DtFloat as i32, ..empty_tensor_proto() };
196        assert!(Tensor::try_from(&no_shape).is_err());
197
198        // -1 is legal in TF ("unknown" dim) but -2 is not, and neither may be
199        // cast to `usize` unchecked: -1 becomes usize::MAX.
200        let negative_dim = TensorProto {
201            dtype: DataType::DtFloat as i32,
202            tensor_shape: Some(TensorShapeProto {
203                dim: vec![Dim { size: -2, name: String::new() }],
204                unknown_rank: false,
205            }),
206            ..empty_tensor_proto()
207        };
208        assert!(Tensor::try_from(&negative_dim).is_err());
209    }
210}
tract_tensorflow/tensor.rs

tract_tensorflow/
tensor.rs
