mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-24 18:18:12 -06:00
434 lines
18 KiB
Rust
434 lines
18 KiB
Rust
//! Model loading, a shared installation self-test, and deterministic test tiles.
|
||
//! `CpuRunner` is the pure-Rust executor; scalar operator tests live in `reference`.
|
||
|
||
use std::path::Path;
|
||
|
||
use web_time::Instant;
|
||
|
||
use serde::Serialize;
|
||
|
||
use crate::manifest::{DenoiserManifest, Gain};
|
||
use crate::run::{Error, TileRunner};
|
||
|
||
/// What `self_test` found.
|
||
#[derive(Clone, Debug, PartialEq, Serialize)]
|
||
#[serde(rename_all = "camelCase")]
|
||
pub struct SelfTest {
|
||
pub ok: bool,
|
||
/// Milliseconds to load and optimise the model.
|
||
pub load_ms: f64,
|
||
/// Median milliseconds to run one tile.
|
||
pub tile_ms: f64,
|
||
/// Megapixels of picture one tile covers, for estimating a whole photo's time.
|
||
pub tile_megapixels: f64,
|
||
/// The model's output scale relative to its input on the test picture (what `Gain::MatchMean` takes out).
|
||
pub scale: f64,
|
||
/// Each check: what was checked and whether it held.
|
||
pub checks: Vec<(String, bool)>,
|
||
}
|
||
|
||
/// The default pure-Rust CPU model, with the same manifest and installation self-test as the reference.
|
||
#[derive(Clone)]
|
||
pub struct CpuRunner {
|
||
runner: crate::cpu::NetRunner,
|
||
manifest: DenoiserManifest,
|
||
load_ms: f64,
|
||
}
|
||
impl CpuRunner {
|
||
pub fn load(path: &Path, manifest: &DenoiserManifest) -> Result<Self, Error> {
|
||
crate::manifest::validate(manifest).map_err(|e| Error::Load(e.to_string()))?;
|
||
let started = Instant::now();
|
||
let net = crate::onnx::read(path).map_err(|e| Error::Load(e.to_string()))?;
|
||
let runner = crate::cpu::NetRunner::new(&net, manifest.tile as usize).map_err(|e| Error::Load(e.to_string()))?;
|
||
Ok(Self { runner, manifest: manifest.clone(), load_ms: started.elapsed().as_secs_f64() * 1e3 })
|
||
}
|
||
pub fn manifest(&self) -> &DenoiserManifest {
|
||
&self.manifest
|
||
}
|
||
pub fn net(&self) -> &crate::net::Net {
|
||
self.runner.net()
|
||
}
|
||
pub fn bytes_per_tile(&self) -> usize {
|
||
self.runner.bytes_per_tile()
|
||
}
|
||
pub fn self_test(&self) -> SelfTest {
|
||
self_test(self, &self.manifest, self.load_ms)
|
||
}
|
||
}
|
||
impl TileRunner for CpuRunner {
|
||
fn run(&self, input: &[f32]) -> Result<Vec<f32>, Error> {
|
||
self.runner.run(input)
|
||
}
|
||
}
|
||
fn self_test(runner: &dyn TileRunner, manifest: &DenoiserManifest, load_ms: f64) -> SelfTest {
|
||
let t = manifest.tile as usize;
|
||
let mut checks: Vec<(String, bool)> = Vec::new();
|
||
let (input, base) = test_tile(t);
|
||
let (mut times, mut scale) = (Vec::new(), 0.0f64);
|
||
let mut run = || {
|
||
let started = Instant::now();
|
||
let r = runner.run(&input);
|
||
times.push(started.elapsed().as_secs_f64() * 1000.0);
|
||
r
|
||
};
|
||
let (a, b) = (run(), run());
|
||
match (&a, &b) {
|
||
(Ok(a), Ok(b)) => {
|
||
checks.push((format!("gives {} finite numbers per tile", a.len()), true));
|
||
checks.push((
|
||
"gives the same answer for the same tile".into(),
|
||
a.iter().zip(b).all(|(x, y)| (x - y).abs() <= 1e-4 * x.abs().max(y.abs()).max(1e-9)),
|
||
));
|
||
let side = 2 * t;
|
||
let mean_in = base.iter().map(|&v| f64::from(v)).sum::<f64>() / base.len().max(1) as f64;
|
||
let mean_out = a.iter().map(|&v| f64::from(v)).sum::<f64>() / a.len().max(1) as f64;
|
||
scale = if mean_in > 0.0 { mean_out / mean_in } else { 0.0 };
|
||
if let Gain::MatchMean { nominal, max_deviation } = manifest.gain {
|
||
let ratio = scale / f64::from(nominal);
|
||
let tolerance = f64::from(max_deviation).max(0.05) * 4.0;
|
||
checks.push((format!("its scale is the {nominal:.3e} the description says (found {scale:.3e})"), (ratio - 1.0).abs() <= tolerance));
|
||
}
|
||
// noise left after bringing the answer back to the input's scale, against the noise that went in
|
||
let g = if scale > 0.0 { 1.0 / scale } else { 1.0 };
|
||
let (mut noise_in, mut noise_out, mut n_in, mut n_out) = (0.0f64, 0.0f64, 0usize, 0usize);
|
||
for (i, (&v, &s)) in input.iter().zip(base.iter().cycle()).enumerate() {
|
||
// centre of each plane only: edges are the model's guess
|
||
let (x, y) = ((i % (t * t)) % t, (i % (t * t)) / t);
|
||
if x >= t / 4 && x < 3 * t / 4 && y >= t / 4 && y < 3 * t / 4 {
|
||
noise_in += f64::from(v - s).powi(2);
|
||
n_in += 1;
|
||
}
|
||
}
|
||
for c in 0..3 {
|
||
for y in side / 4..3 * side / 4 {
|
||
for x in side / 4..3 * side / 4 {
|
||
if let (Some(&o), Some(&s)) = (a.get(c * side * side + y * side + x), base.get((y / 2) * t + x / 2)) {
|
||
noise_out += (f64::from(o) * g - f64::from(s)).powi(2);
|
||
n_out += 1;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let (rms_in, rms_out) = ((noise_in / n_in.max(1) as f64).sqrt(), (noise_out / n_out.max(1) as f64).sqrt());
|
||
checks.push((format!("takes noise down (from {rms_in:.4} to {rms_out:.4})"), rms_out < 0.8 * rms_in));
|
||
}
|
||
(Err(e), _) | (_, Err(e)) => checks.push((format!("runs: {e}"), false)),
|
||
}
|
||
times.sort_by(|x, y| x.total_cmp(y));
|
||
SelfTest {
|
||
ok: !checks.is_empty() && checks.iter().all(|(_, ok)| *ok),
|
||
load_ms,
|
||
tile_ms: times.get(times.len() / 2).copied().unwrap_or(0.0),
|
||
tile_megapixels: (2 * t * 2 * t) as f64 / 1e6,
|
||
scale,
|
||
checks,
|
||
}
|
||
}
|
||
|
||
/// A deterministic test tile: four planes of a smooth gradient with noise on top, and the gradient alone for one
|
||
/// plane (`tile × tile`).
|
||
pub fn test_tile(t: usize) -> (Vec<f32>, Vec<f32>) {
|
||
if t == 0 || t > 2048 {
|
||
return (Vec::new(), Vec::new());
|
||
}
|
||
let mut seed = 0x2545_f491u32;
|
||
let mut next = move || {
|
||
seed ^= seed << 13;
|
||
seed ^= seed >> 17;
|
||
seed ^= seed << 5;
|
||
(seed >> 8) as f32 / (1u32 << 24) as f32
|
||
};
|
||
let base: Vec<f32> = (0..t * t).map(|i| 0.12 + 0.3 * (((i % t) + (i / t)) as f32 / (2 * t) as f32)).collect();
|
||
let input = (0..4 * t * t)
|
||
.map(|i| {
|
||
// sum of four uniforms: close to Gaussian, deviation 0.0577 × 2 before scaling
|
||
let n = next() + next() + next() + next() - 2.0;
|
||
base.get(i % (t * t)).copied().unwrap_or(0.0) + 0.05 * n
|
||
})
|
||
.collect();
|
||
(input, base)
|
||
}
|
||
|
||
/// A tile for checking that two runners of one model agree: the whole range a raw photo has rather than the test tile's
|
||
/// mid-greys — a noisy ramp from black to white across the tile, a flat black corner and a clipped white one (four
|
||
/// planes of `tile × tile`).
|
||
pub fn check_tile(t: usize) -> Vec<f32> {
|
||
if t == 0 || t > 2048 {
|
||
return Vec::new();
|
||
}
|
||
let mut seed = 0x9e37_79b9u32;
|
||
let mut next = move || {
|
||
seed ^= seed << 13;
|
||
seed ^= seed >> 17;
|
||
seed ^= seed << 5;
|
||
(seed >> 8) as f32 / (1u32 << 24) as f32
|
||
};
|
||
let (edge, span) = ((t / 8).max(1), (2 * t).saturating_sub(2).max(1) as f32);
|
||
(0..4 * t * t)
|
||
.map(|i| {
|
||
let (x, y) = ((i % (t * t)) % t, (i % (t * t)) / t);
|
||
let n = next() + next() + next() + next() - 2.0;
|
||
if x < edge && y < edge {
|
||
0.0
|
||
} else if x + edge >= t && y + edge >= t {
|
||
1.0
|
||
} else {
|
||
((x + y) as f32 / span + 0.05 * n).clamp(0.0, 1.0)
|
||
}
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::licence::Licence;
|
||
use crate::manifest::Domain;
|
||
use crate::synthetic_proto::{len_field, value_info_bytes};
|
||
|
||
fn varint(mut v: u64, out: &mut Vec<u8>) {
|
||
loop {
|
||
let b = (v & 0x7f) as u8;
|
||
v >>= 7;
|
||
if v == 0 {
|
||
out.push(b);
|
||
return;
|
||
}
|
||
out.push(b | 0x80);
|
||
}
|
||
}
|
||
|
||
fn varint_field(n: u64, v: u64, out: &mut Vec<u8>) {
|
||
varint(n << 3, out);
|
||
varint(v, out);
|
||
}
|
||
|
||
fn node(op: &str, inputs: &[&str], output: &str, ints: &[(&str, i64)]) -> Vec<u8> {
|
||
let mut n = Vec::new();
|
||
for i in inputs {
|
||
len_field(1, i.as_bytes(), &mut n);
|
||
}
|
||
len_field(2, output.as_bytes(), &mut n);
|
||
len_field(4, op.as_bytes(), &mut n);
|
||
for (name, value) in ints {
|
||
let mut a = Vec::new();
|
||
len_field(1, name.as_bytes(), &mut a);
|
||
varint_field(3, *value as u64, &mut a);
|
||
varint_field(20, 2, &mut a);
|
||
len_field(5, &a, &mut n);
|
||
}
|
||
n
|
||
}
|
||
|
||
fn tensor(name: &str, dims: &[u64], values: &[f32]) -> Vec<u8> {
|
||
let mut t = Vec::new();
|
||
for d in dims {
|
||
varint_field(1, *d, &mut t);
|
||
}
|
||
varint_field(2, 1, &mut t);
|
||
len_field(8, name.as_bytes(), &mut t);
|
||
let raw: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||
len_field(9, &raw, &mut t);
|
||
t
|
||
}
|
||
|
||
/// A tiny working denoiser-shaped network: `[1, 4, t, t]` → a 1 × 1 convolution to 12 channels (each output
|
||
/// channel is `scale` times one input plane, so the scale is known) → depth-to-space → `[1, 3, 2t, 2t]`.
|
||
fn tiny_model(t: u64, scale: f32) -> Vec<u8> {
|
||
tiny_model_with(t, scale, 12)
|
||
}
|
||
|
||
/// The same with another number of convolution channels; only 12 can be arranged into the three colours.
|
||
fn tiny_model_with(t: u64, scale: f32, channels: u64) -> Vec<u8> {
|
||
let mut graph = Vec::new();
|
||
len_field(1, &node("Conv", &["data", "w"], if channels == 12 { "c" } else { "rgb" }, &[("group", 1)]), &mut graph);
|
||
if channels == 12 {
|
||
let mut depth_to_space = node("DepthToSpace", &["c"], "rgb", &[("blocksize", 2)]);
|
||
let mut mode = Vec::new();
|
||
len_field(1, b"mode", &mut mode);
|
||
len_field(4, b"CRD", &mut mode);
|
||
varint_field(20, 3, &mut mode);
|
||
len_field(5, &mode, &mut depth_to_space);
|
||
len_field(1, &depth_to_space, &mut graph);
|
||
}
|
||
// CRD output channel colour·4 + dy·2 + dx reads input plane [R, G1, B] of the tile.
|
||
let mut w = vec![0f32; channels as usize * 4];
|
||
for oc in 0..channels as usize {
|
||
let colour = (oc / 4).min(2);
|
||
let plane = [0, 1, 3][colour];
|
||
w[oc * 4 + plane] = scale;
|
||
}
|
||
len_field(5, &tensor("w", &[channels, 4, 1, 1], &w), &mut graph);
|
||
len_field(11, &value_info_bytes("data", 1, &[Ok(1), Ok(4), Ok(t), Ok(t)]), &mut graph);
|
||
len_field(12, &value_info_bytes("rgb", 1, &[Ok(1), Ok(3), Ok(2 * t), Ok(2 * t)]), &mut graph);
|
||
let mut ops = Vec::new();
|
||
len_field(1, b"", &mut ops);
|
||
varint_field(2, 13, &mut ops);
|
||
let mut model = Vec::new();
|
||
varint_field(1, 8, &mut model);
|
||
len_field(7, &graph, &mut model);
|
||
len_field(8, &ops, &mut model);
|
||
model
|
||
}
|
||
|
||
fn manifest(tile: u32, nominal: f32) -> DenoiserManifest {
|
||
DenoiserManifest {
|
||
id: "tiny".into(),
|
||
name: "Tiny".into(),
|
||
version: "1".into(),
|
||
licence: Licence::default(),
|
||
source: None,
|
||
sha256: None,
|
||
size_bytes: None,
|
||
provenance: String::new(),
|
||
domain: Domain::BayerToRgb,
|
||
tile,
|
||
overlap: 16,
|
||
gain: Gain::MatchMean { nominal, max_deviation: 0.05 },
|
||
}
|
||
}
|
||
|
||
fn temp_model(name: &str, bytes: &[u8]) -> std::path::PathBuf {
|
||
let p = std::env::temp_dir().join(format!("lc-denoise-runtime-{name}-{}.onnx", std::process::id()));
|
||
std::fs::write(&p, bytes).unwrap();
|
||
p
|
||
}
|
||
|
||
#[test]
|
||
fn a_working_model_loads_runs_and_has_the_scale_it_says() {
|
||
let p = temp_model("ok", &tiny_model(64, 1000.0));
|
||
let r = CpuRunner::load(&p, &manifest(64, 1000.0)).unwrap();
|
||
let input: Vec<f32> = (0..4 * 64 * 64).map(|i| (i % 97) as f32 / 97.0).collect();
|
||
let out = r.run(&input).unwrap();
|
||
assert_eq!(out.len(), 3 * 128 * 128);
|
||
// output (colour c, pixel (2y+dy, 2x+dx)) = scale × input plane [R, G1, B][c] at (y, x)
|
||
for (c, plane) in [0usize, 1, 3].into_iter().enumerate() {
|
||
let (y, x) = (5usize, 9usize);
|
||
assert!((out[c * 128 * 128 + (2 * y + 1) * 128 + 2 * x] - 1000.0 * input[plane * 64 * 64 + y * 64 + x]).abs() < 1e-2);
|
||
}
|
||
let _ = std::fs::remove_file(p);
|
||
}
|
||
|
||
#[test]
|
||
fn the_self_test_checks_scale_and_noise() {
|
||
// a model that only multiplies cannot take noise down: it must fail that check but pass the others
|
||
let p = temp_model("selftest", &tiny_model(64, 1000.0));
|
||
let t = CpuRunner::load(&p, &manifest(64, 1000.0)).unwrap().self_test();
|
||
assert!(!t.ok, "{t:?}");
|
||
assert!(t.checks.iter().any(|(c, ok)| c.contains("same answer") && *ok), "{t:?}");
|
||
assert!(t.checks.iter().any(|(c, ok)| c.contains("scale") && *ok), "{t:?}");
|
||
assert!(t.checks.iter().any(|(c, ok)| c.contains("noise") && !*ok), "{t:?}");
|
||
// and one that claims the wrong scale fails that check
|
||
let q = CpuRunner::load(&p, &manifest(64, 5.0)).unwrap().self_test();
|
||
assert!(q.checks.iter().any(|(c, ok)| c.contains("scale") && !*ok), "{q:?}");
|
||
let _ = std::fs::remove_file(p);
|
||
}
|
||
|
||
#[test]
|
||
fn wrong_sizes_and_broken_files_are_errors() {
|
||
let p = temp_model("sizes", &tiny_model(64, 1.0));
|
||
let r = CpuRunner::load(&p, &manifest(64, 1.0)).unwrap();
|
||
assert!(matches!(r.run(&[0.0; 10]), Err(Error::Input(_))));
|
||
// a fully convolutional model works at whatever tile size the manifest asks for
|
||
let wide = CpuRunner::load(&p, &manifest(128, 1.0)).unwrap();
|
||
assert_eq!(wide.run(&vec![0.5; 4 * 128 * 128]).unwrap().len(), 3 * 256 * 256);
|
||
// a model that gives another number of values than the manifest promises is an error
|
||
let q = temp_model("channels", &tiny_model_with(64, 1.0, 8));
|
||
assert!(CpuRunner::load(&q, &manifest(64, 1.0)).is_err(), "shape is refused before execution");
|
||
let _ = std::fs::remove_file(q);
|
||
for (name, bytes) in [("junk", b"nonsense".to_vec()), ("empty", Vec::new())] {
|
||
let q = temp_model(name, &bytes);
|
||
assert!(CpuRunner::load(&q, &manifest(64, 1.0)).is_err(), "{name}");
|
||
let _ = std::fs::remove_file(q);
|
||
}
|
||
assert!(CpuRunner::load(&std::env::temp_dir().join("does-not-exist.onnx"), &manifest(64, 1.0)).is_err());
|
||
let _ = std::fs::remove_file(p);
|
||
}
|
||
|
||
#[test]
|
||
fn the_check_tile_covers_black_to_white() {
|
||
for t in [1, 2, 7, 64] {
|
||
let c = check_tile(t);
|
||
assert_eq!(c.len(), 4 * t * t);
|
||
assert!(c.iter().all(|v| (0.0..=1.0).contains(v)), "{t}");
|
||
}
|
||
let c = check_tile(64);
|
||
let near = |lo: f32, hi: f32| c.iter().filter(|v| (lo..=hi).contains(*v)).count();
|
||
assert!(near(0.0, 0.0) >= 4 * 8 * 8 && near(1.0, 1.0) >= 4 * 8 * 8, "a flat black and a clipped white corner");
|
||
for band in 0..10 {
|
||
let lo = band as f32 / 10.0;
|
||
assert!(near(lo, lo + 0.1) > 100, "values around {lo}");
|
||
}
|
||
assert_eq!(check_tile(64), c, "the same tile every time");
|
||
}
|
||
|
||
#[test]
|
||
fn truncated_and_corrupted_models_never_panic() {
|
||
let good = tiny_model(64, 1.0);
|
||
for n in (0..good.len()).step_by(5) {
|
||
let q = temp_model("trunc", &good[..n]);
|
||
let _ = CpuRunner::load(&q, &manifest(64, 1.0));
|
||
let _ = std::fs::remove_file(q);
|
||
}
|
||
for at in (0..good.len()).step_by(7) {
|
||
let mut bad = good.clone();
|
||
if let Some(b) = bad.get_mut(at) {
|
||
*b = b.wrapping_add(0x55);
|
||
}
|
||
let q = temp_model("flip", &bad);
|
||
if let Ok(r) = CpuRunner::load(&q, &manifest(64, 1.0)) {
|
||
let _ = r.run(&vec![0.5; 4 * 64 * 64]);
|
||
}
|
||
let _ = std::fs::remove_file(q);
|
||
}
|
||
}
|
||
|
||
/// Opt-in: `LC_DENOISE_MODEL=<model_bayer.onnx> cargo test -p lightcraft-denoise --release
|
||
/// -- --ignored --nocapture` self-tests a real model (RawNIND's Bayer model: tile 512, scale about 1e6).
|
||
#[test]
|
||
#[ignore = "needs a real model file: set LC_DENOISE_MODEL"]
|
||
fn a_real_model_passes_the_self_test() {
|
||
let Some(path) = std::env::var_os("LC_DENOISE_MODEL") else { return };
|
||
let m = manifest(512, 1.0e6);
|
||
let r = CpuRunner::load(Path::new(&path), &m).unwrap();
|
||
let t = r.self_test();
|
||
println!("load {:.0} ms | {:.0} ms per tile | scale {:.4e} | {:?}", t.load_ms, t.tile_ms, t.scale, t.checks);
|
||
assert!(t.ok);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod cpu_tests {
|
||
use super::*;
|
||
#[test]
|
||
fn manifest_and_test_tile_sizes_are_checked_without_tract() {
|
||
for t in [0, 2049, usize::MAX] {
|
||
assert!(test_tile(t).0.is_empty());
|
||
assert!(check_tile(t).is_empty());
|
||
}
|
||
let path = std::env::temp_dir().join(format!("lc-cpu-load-{}.onnx", std::process::id()));
|
||
std::fs::write(&path, crate::synthetic::unet_onnx(64, 5, 1, 7)).unwrap();
|
||
let mut m = DenoiserManifest {
|
||
id: "synthetic".into(),
|
||
name: "Synthetic".into(),
|
||
version: "1".into(),
|
||
licence: Default::default(),
|
||
source: None,
|
||
sha256: None,
|
||
size_bytes: None,
|
||
provenance: "synthetic test".into(),
|
||
domain: crate::manifest::Domain::BayerToRgb,
|
||
tile: 64,
|
||
overlap: 16,
|
||
gain: Gain::None,
|
||
};
|
||
let r = CpuRunner::load(&path, &m).unwrap();
|
||
let x = check_tile(64);
|
||
assert_eq!(r.run(&x).unwrap(), r.run(&x).unwrap());
|
||
assert!(r.self_test().checks.iter().any(|(text, ok)| text.contains("same answer") && *ok));
|
||
m.tile = u32::MAX;
|
||
assert!(CpuRunner::load(&path, &m).is_err());
|
||
std::fs::remove_file(path).unwrap();
|
||
}
|
||
}
|