Merge pull request #20 from bminixhofer/symbolic-dim

Support for symbolic dims

Author: bminixhofer

Cargo.lock

@@ -10,10 +10,10 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 wasm-bindgen = "^0.2"
 js-sys = "0.3.39"
-tract-core = "0.9.2"
-tract-onnx = "0.9.2"
-tract-hir = "0.9.2"
-tract-tensorflow = "0.9.2"
+tract-core = "0.11.1"
+tract-onnx = "0.11.1"
+tract-hir = "0.11.1"
+tract-tensorflow = "0.11.1"
 console_error_panic_hook = "0.1.1"
 ndarray-rand = { version = "0.11.0", optional = true }
 serde = { version = "1.0", features = ["derive"], optional = true }

Cargo.toml

@@ -80,19 +80,29 @@ tractjs.load("./path/to/your/model").then((model) => {
 
 ### My model with dynamic input dimensions doesn't work. Why?
 
-Currently, tract requires fully determined input dimensions to optimize a model. There are two options:
+Currently, tract requires has some restrictions on dynamic dimensions. If your model has a dynamic dimension, there's multiple solutions:
 
-1. Turn `optimize` off:
+1. Declare a dynamic dimension via an input fact. Input facts are a way to provide additional information about input type and shape that can not be inferred via the model data:
 
 ```js
 const model = await tractjs.load("path/to/your/model", {
-  optimize: false,
+  inputFacts: {
+    0: ["float32", [1, "s", 224, 224]],
+  },
 });
 ```
 
-This will however _significantly_ impact performance.
+This dimension must then be made concrete on prediction:
+
+```js
+model.predict(input, {
+  "s": 3 // or some other value
+})
+```
+
+The API supports multiple dynamic dimensions, but currently it will probably only work with one.
 
-2. Set fixed input dimensions via input facts. Input facts are a way to provide additional information about input type and shape that can not be inferred via the model data:
+2. Set fixed input dimensions via input facts. This is of course not ideal because subsequently the model can only be passed inputs with this exact shape:
 
 ```js
 const model = await tractjs.load("path/to/your/model", {
@@ -104,13 +114,17 @@ const model = await tractjs.load("path/to/your/model", {
 });
 ```
 
-Be aware that the model will only work properly with inputs of this exact shape though.
+3. Turn `optimize` off. This is the nuclear option. It will turn off all optimizations relying on information about input shape. This will make sure your model work (even with multiple dynamic dimensions) but _significantly_ impact performance:
 
-There is ongoing work in tract to allow dynamically sized inputs.
+```js
+const model = await tractjs.load("path/to/your/model", {
+  optimize: false,
+});
+```
 
 ### What about size?
 
-At the time of writing, tractjs is very large for web standards (8.5MB raw, 2.5MB gzipped). This is due to tract being quite large, and due to some overhead from inlining the WASM. But it's not as bad as it sounds. You can load tractjs lazily along your demo, where you will likely have to load significantly large weights too.
+At the time of writing, tractjs is very large for web standards (6.2MB raw, 2.1MB gzipped). This is due to tract being quite large, and due to some overhead from inlining the WASM. But it's not as bad as it sounds. You can load tractjs lazily along your demo, where you will likely have to load significantly large weights too.
 
 If you are working on a very size-sensitive application, get in touch and we can work on decreasing the size. There are some more optimizations to be done (e. g. an option not to inline WASM, and removing panics from the build). There is also ongoing work in tract to decrease size.
 

README.md

@@ -125,8 +125,7 @@ fn custom_input_tf<P1: AsRef<Path>, P2: AsRef<Path>>(
     input_file: P1,
     output_file: P2,
 ) -> TractResult<()> {
-    let mut cursor =
-        Cursor::new(include_bytes!("../models/data/squeezenet_1_1/model.pb") as &[u8]);
+    let mut cursor = Cursor::new(include_bytes!("../models/data/squeezenet_1_1/model.pb") as &[u8]);
     let model = tract_tensorflow::tensorflow()
         .model_for_read(&mut cursor)?
         .with_input_fact(
@@ -137,7 +136,10 @@ fn custom_input_tf<P1: AsRef<Path>, P2: AsRef<Path>>(
         .into_optimized()?
         .into_runnable()?;
 
-    let inputs = tvec![random_from_shape((1, 27, 27, 128)), random_from_shape((1, 27, 27, 128))];
+    let inputs = tvec![
+        random_from_shape((1, 27, 27, 128)),
+        random_from_shape((1, 27, 27, 128))
+    ];
     serialize_tensors(&inputs, input_file);
 
     let preds = model.run(inputs)?;

quality/tests/reference.rs

@@ -1,5 +1,6 @@
-use js_sys::{Array, Error, Object, Uint32Array};
+use js_sys::{Array, Error, JsString, Object, Uint32Array};
 use std::io::Cursor;
+use tract_core::internal::ToDim;
 use tract_hir::prelude::*;
 use tract_onnx::prelude::*;
 use tract_tensorflow::prelude::*;
@@ -31,8 +32,9 @@ impl<T: std::fmt::Debug> TractResultExt<T> for TractResult<T> {
     }
 }
 
-fn fact_from_js(input: JsValue) -> InferenceFact {
-    let input: Array = input.dyn_into().expect("fact must be an Array.");
+fn fact_from_js(input: JsValue) -> (InferenceFact, bool) {
+    let mut has_symbolic_dim = false;
+    let input: Array = input.dyn_into().expect("Fact must be an Array.");
     let dtype = if let Some(string) = input.get(0).as_string() {
         Some(match string.as_str() {
             "int8" => i8::datum_type(),
@@ -52,19 +54,77 @@ fn fact_from_js(input: JsValue) -> InferenceFact {
         Some(
             shape
                 .iter()
-                .map(|x| x.as_f64().expect("fact[1] must be an Array of numbers.") as usize)
+                .map(|x| {
+                    let mut dim: Option<TDim> = None;
+
+                    if let Some(number) = x.as_f64() {
+                        dim = Some((number as isize).to_dim());
+                    } else if let Some(string) = x.as_string() {
+                        has_symbolic_dim = true;
+                        dim = Some(Symbol::from(string.chars().next().expect("fact[1][i] must consist of exactly one character if string.")).into());
+                    } else if let Ok(object) = x.dyn_into::<Object>() {
+                        // could be made prettier - maybe with serde, but would add an additional dependency
+                        let id = js_sys::Reflect::get(&object, &JsString::from("id")).expect("fact[1][i] must have 'id' key if object.")
+                            .as_string().expect("fact[1][i]['id'] must be a string.");
+
+                        let intercept = js_sys::Reflect::get(&object, &JsString::from("intercept")).expect("fact[1][i] must have 'intercept' key if object.")
+                            .as_f64().expect("fact[1][i]['intercept'] must be a number.") as isize;
+
+                        let slope = js_sys::Reflect::get(&object, &JsString::from("slope")).expect("fact[1][i] must have 'slope' key if object.")
+                            .as_f64().expect("fact[1][i]['slope'] must be a number.");
+
+                        let symbol: TDim =  Symbol::from(id.chars().next().expect("fact[1][i] must consist of exactly one character if string.")).into();
+                        has_symbolic_dim = true;
+
+                        dim = Some((symbol * slope) + intercept)
+                    }
+
+                    dim.expect("fact[1][i] must be one of: number, string, or object with slope, intercept and id keys.")
+                })
                 .collect::<Vec<_>>(),
         )
     } else {
         None
     };
 
-    match (dtype, shape) {
-        (Some(dtype), Some(shape)) => InferenceFact::dt_shape(dtype, shape),
-        (Some(dtype), None) => InferenceFact::dt(dtype),
-        (None, Some(shape)) => InferenceFact::shape(shape),
-        (None, None) => panic!("either dtype or shape must be specified."),
+    (
+        match (dtype, shape) {
+            (Some(dtype), Some(shape)) => InferenceFact::dt_shape(dtype, shape),
+            (Some(dtype), None) => InferenceFact::dt(dtype),
+            (None, Some(shape)) => InferenceFact::shape(shape),
+            (None, None) => panic!("either dtype or shape must be specified."),
+        },
+        has_symbolic_dim,
+    )
+}
+
+fn symbol_values_from_js(input: JsValue) -> SymbolValues {
+    let input: Object = input.dyn_into().expect("SymbolValues must be object.");
+    let entries = Object::entries(&input);
+
+    let mut symbol_values = SymbolValues::default();
+
+    for i in 0..entries.length() {
+        let val: Array = entries
+            .get(i)
+            .dyn_into()
+            .expect(".entries[i] must be an array.");
+
+        symbol_values = symbol_values.with(
+            val.get(0)
+                .as_string()
+                .expect("entries[i][0] must be a string.")
+                .chars()
+                .next()
+                .expect("entries[i][0] must consist of exactly one character.")
+                .into(),
+            val.get(1)
+                .as_f64()
+                .expect("entries[i][1] must be a number.") as i64,
+        )
     }
+
+    symbol_values
 }
 
 #[wasm_bindgen]
@@ -189,14 +249,33 @@ impl From<CoreTensor> for Tensor {
 
 enum Model {
     Inference(tract_hir::infer::InferenceSimplePlan<InferenceModel>),
+    Optimized(TypedModel),
     Typed(TypedSimplePlan<TypedModel>),
 }
 
+#[wasm_bindgen]
+extern "C" {
+    #[wasm_bindgen(js_namespace = console)]
+    pub fn log(input: &str);
+}
+
 impl Model {
-    fn run(&self, inputs: TVec<Tensor>) -> TractResult<TVec<Arc<Tensor>>> {
+    fn run(
+        &self,
+        inputs: TVec<Tensor>,
+        symbol_values: &SymbolValues,
+    ) -> TractResult<TVec<Arc<Tensor>>> {
         match self {
             Model::Inference(x) => x.run(inputs),
             Model::Typed(x) => x.run(inputs),
+            Model::Optimized(x) => {
+                let model = x
+                    .concretize_dims(symbol_values)?
+                    .optimize()?
+                    .into_runnable()?;
+
+                model.run(inputs)
+            }
         }
     }
 }
@@ -207,6 +286,12 @@ impl From<tract_hir::infer::InferenceSimplePlan<InferenceModel>> for Model {
     }
 }
 
+impl From<TypedModel> for Model {
+    fn from(input: TypedModel) -> Self {
+        Model::Optimized(input)
+    }
+}
+
 impl From<TypedSimplePlan<TypedModel>> for Model {
     fn from(input: TypedSimplePlan<TypedModel>) -> Self {
         Model::Typed(input)
@@ -231,6 +316,7 @@ impl CoreModel {
         console_error_panic_hook::set_once();
 
         let mut reader = Cursor::new(data);
+        let mut model_has_symbolic_dim = false;
 
         let mut model = if use_onnx {
             onnx().model_for_read(&mut reader)
@@ -243,7 +329,8 @@ impl CoreModel {
             .iter()
             .zip(Object::values(&input_facts).iter())
         {
-            let fact = fact_from_js(fact);
+            let (fact, fact_has_symbolic_dim) = fact_from_js(fact);
+            model_has_symbolic_dim = model_has_symbolic_dim || fact_has_symbolic_dim;
 
             model
                 .set_input_fact(
@@ -275,22 +362,31 @@ impl CoreModel {
                 .map_js_error()?;
         }
 
-        let model: Model = if optimize {
+        let model: Model = if optimize && !model_has_symbolic_dim {
             model
                 .into_optimized()
                 .map_js_error()?
                 .into_runnable()
                 .map_js_error()?
                 .into()
+        } else if optimize {
+            model.into_optimized().map_js_error()?.into()
         } else {
             model.into_runnable().map_js_error()?.into()
         };
 
         Ok(CoreModel { model })
     }
 
-    pub fn predict(&self, data: CoreTensorVec) -> Result<CoreTensorVec, JsValue> {
-        let outputs = self.model.run(data.into_tvec()).map_js_error()?;
+    pub fn predict(
+        &self,
+        data: CoreTensorVec,
+        symbol_values: JsValue,
+    ) -> Result<CoreTensorVec, JsValue> {
+        let outputs = self
+            .model
+            .run(data.into_tvec(), &symbol_values_from_js(symbol_values))
+            .map_js_error()?;
         Ok(CoreTensorVec::from_slice(&outputs))
     }
 }