use ember_infer_macros::model;
use ember_infer_ref::RefBackend;
extern crate self as microflow;
extern crate self as nalgebra;
pub mod buffer {
pub type Buffer2D<T, const ROWS: usize, const COLS: usize> = [[T; COLS]; ROWS];
pub type Buffer4D<
T,
const BATCH: usize,
const HEIGHT: usize,
const WIDTH: usize,
const CHANNELS: usize,
> = [[[[T; CHANNELS]; WIDTH]; HEIGHT]; BATCH];
}
#[macro_export]
macro_rules! matrix {
($( $($value:expr),+ );+ $(;)?) => {
[$( [ $($value),+ ] ),+]
};
}
#[allow(unused_imports)]
#[path = "fixtures/speech.rs"]
mod speech_samples;
#[allow(unused_imports)]
#[path = "fixtures/person_detect.rs"]
mod person_detect_samples;
#[model("models/sine.tflite")]
struct SineModel;
#[model("models/speech.tflite")]
struct SpeechModel;
#[model("models/person_detect.tflite")]
struct PersonDetectModel;
fn quantize(val: f32, scale: f32, zero_point: i32) -> i8 {
((val / scale) + zero_point as f32)
.round()
.clamp(i8::MIN as f32, i8::MAX as f32) as i8
}
fn dequantize(val: i8, scale: f32, zero_point: i32) -> f32 {
(val as i32 - zero_point) as f32 * scale
}
fn assert_top_class(model: &str, sample: &str, scores: &[f32], expected: usize) {
let (actual, actual_score) = scores
.iter()
.copied()
.enumerate()
.max_by(|(_, left), (_, right)| left.total_cmp(right))
.unwrap();
assert_eq!(
actual, expected,
"{model} {sample} expected class {expected}, got class {actual}; scores={scores:?}"
);
assert!(
actual_score > 0.5,
"{model} {sample} expected class {expected} confidence > 0.5, got {actual_score}; scores={scores:?}"
);
}
fn run_speech_sample(name: &str, input: &[i8]) -> Vec<f32> {
assert_eq!(input.len(), SpeechModel::input_len());
let mut backend = RefBackend;
let mut output = vec![0i8; SpeechModel::output_len()];
SpeechModel::predict_quantized(&mut backend, input, &mut output).unwrap();
let scores: Vec<f32> = output
.iter()
.map(|&v| {
dequantize(
v,
SpeechModel::output_scale(),
SpeechModel::output_zero_point(),
)
})
.collect();
println!("[speech:{name}] input length = {}", input.len());
println!("[speech:{name}] output length = {} classes", output.len());
for (i, score) in scores.iter().enumerate() {
println!("[speech:{name}] class {i} score = {score:.4}");
}
let sum: f32 = scores.iter().sum();
println!("[speech:{name}] sum of scores = {sum:.4} (should be ~= 1.0)");
assert_eq!(output.len(), SpeechModel::output_len());
for (index, score) in scores.iter().copied().enumerate() {
assert!(
(0.0..=1.0).contains(&score),
"speech {name} output[{index}] out of range: score={score}"
);
}
assert!(
(sum - 1.0).abs() < 0.05,
"speech {name} scores should sum to ~= 1.0, got {sum}"
);
scores
}
fn run_person_detect_sample(name: &str, input: &[i8]) -> Vec<f32> {
assert_eq!(input.len(), PersonDetectModel::input_len());
let mut backend = RefBackend;
let mut output = vec![0i8; PersonDetectModel::output_len()];
PersonDetectModel::predict_quantized(&mut backend, input, &mut output).unwrap();
println!(
"[person_detect:{name}] input length = {} pixels (i8)",
input.len()
);
println!(
"[person_detect:{name}] output length = {} classes",
output.len()
);
let scores: Vec<f32> = output
.iter()
.enumerate()
.map(|(i, &v)| {
let label = match i {
0 => "no_person",
1 => "person",
_ => "unknown",
};
let score = dequantize(
v,
PersonDetectModel::output_scale(),
PersonDetectModel::output_zero_point(),
);
println!(
"[person_detect:{name}] class {i} ({label}) raw i8 = {v} -> score = {score:.4}"
);
score
})
.collect();
assert_eq!(output.len(), PersonDetectModel::output_len());
for (index, score) in scores.iter().copied().enumerate() {
assert!(
(0.0..=1.0).contains(&score),
"person_detect {name} output[{index}] out of range: score={score}"
);
}
scores
}
fn assert_sine_close(x: f32, expected: f32) {
let mut backend = RefBackend;
let q_input = quantize(x, SineModel::input_scale(), SineModel::input_zero_point());
let input = [q_input; SineModel::input_len()];
let mut output = [0i8; SineModel::output_len()];
SineModel::predict_quantized(&mut backend, &input, &mut output).unwrap();
let result = dequantize(
output[0],
SineModel::output_scale(),
SineModel::output_zero_point(),
);
println!(
"[sine] input x = {:.4} -> quantized input = {}",
x, q_input
);
println!(
"[sine] raw output i8 = {} -> dequantized f32 = {:.4}",
output[0], result
);
println!(
"[sine] expected sin({:.4}) = {:.4} | error = {:.4} | tolerance = 0.10 {}",
x,
expected,
(result - expected).abs(),
if (result - expected).abs() < 0.1 {
"✓"
} else {
"✗"
}
);
assert!(
(result - expected).abs() < 0.1,
"expected sin({x}) ~= {expected}, got {result}; q_input={q_input}, q_output={:?}",
output
);
}
#[test]
fn sine_end_to_end_zero() {
assert_sine_close(0.0, 0.0);
}
#[test]
fn sine_end_to_end_pi_over_2() {
assert_sine_close(core::f32::consts::FRAC_PI_2, 1.0);
}
#[test]
fn sine_end_to_end_pi() {
assert_sine_close(core::f32::consts::PI, 0.0);
}
#[test]
fn speech_compiles_and_runs() {
let yes_scores = run_speech_sample("yes", speech_samples::YES.as_flattened());
let no_scores = run_speech_sample("no", speech_samples::NO.as_flattened());
assert_top_class("speech", "yes", &yes_scores, 2);
assert_top_class("speech", "no", &no_scores, 3);
assert_ne!(
yes_scores, no_scores,
"speech samples should produce different outputs"
);
}
#[test]
fn person_detect_compiles_and_runs() {
let person_scores = run_person_detect_sample(
"person",
person_detect_samples::PERSON
.as_flattened()
.as_flattened()
.as_flattened(),
);
let no_person_scores = run_person_detect_sample(
"no_person",
person_detect_samples::NO_PERSON
.as_flattened()
.as_flattened()
.as_flattened(),
);
assert_top_class("person_detect", "person", &person_scores, 1);
assert_top_class("person_detect", "no_person", &no_person_scores, 0);
assert_ne!(
person_scores, no_person_scores,
"person_detect samples should produce different outputs"
);
}