dtourolle
14f08a565f
A 20 MP frame spent 0.32 s outside the network: each tile's mosaic and sigma gathered on one thread, then its 24 MB output copied out of the runtime and back into the frame, all in series with the device. Tiles are now gathered on every core by a producer thread one tile ahead, so the gather overlaps the run; the centre is written back across cores; and the tile interface hands its inputs over and lends its output, so neither side is copied. With a stand-in network that does nothing, the tiler's own time falls to 0.14 s at the 1408 tile and 0.09 s at 2048. The exactness and Bayer-phase tests are unchanged and pass.
79 lines
2.8 KiB
Rust
79 lines
2.8 KiB
Rust
//! TRACES: FR-DEV-3g
|
||
//! The denoise network under the inference engine.
|
||
//!
|
||
//! The shipped export takes `mosaic` and `sigma`, `1×1×1408×1408`, and
|
||
//! returns `rgb`, `1×3×1408×1408` (darkroom-denoise `denoise/export.py`,
|
||
//! fixed shape because every model the engine runs is). The engine picks the
|
||
//! rung: fp16 on TensorRT and MIGraphX, which measured 0.00 dB from f32; f32
|
||
//! on CUDA and the CPU; on the Hexagon the `.a16w16.onnx` sibling, 16-bit
|
||
//! activations and weights, 0.00 dB from f32 on the tablet itself where int8
|
||
//! lost 5–9 dB (docs/dev/inference.md §1.5). That sibling is the same network
|
||
//! with the Bayer packing spelled `SpaceToDepth`, which QNN can hold and the
|
||
//! 6-D reshape it replaces it cannot.
|
||
|
||
use crate::tile::TileNet;
|
||
use crate::DenoiseError;
|
||
use dr_inference_engine::{Model, Role};
|
||
|
||
/// The edge of the tile the shipped export takes.
|
||
pub const TILE: usize = 1408;
|
||
|
||
pub struct OnnxNet {
|
||
model: Model,
|
||
tile: usize,
|
||
}
|
||
|
||
impl OnnxNet {
|
||
pub fn from_path(path: &std::path::Path) -> Result<Self, DenoiseError> {
|
||
let (path, form) = dr_inference_engine::resolve_model(Role::Denoiser, path);
|
||
let bytes = std::fs::read(&path)?;
|
||
Ok(OnnxNet {
|
||
model: dr_inference_engine::open(Role::Denoiser, form, &bytes)?,
|
||
tile: TILE,
|
||
})
|
||
}
|
||
|
||
/// Where it runs, for a status line.
|
||
pub fn rung(&self) -> Result<dr_inference_engine::Rung, DenoiseError> {
|
||
Ok(self.model.acquire()?.rung())
|
||
}
|
||
}
|
||
|
||
impl TileNet for OnnxNet {
|
||
fn tile(&self) -> usize {
|
||
self.tile
|
||
}
|
||
|
||
fn run(
|
||
&mut self,
|
||
mosaic: Vec<f32>,
|
||
sigma: Vec<f32>,
|
||
write: &mut dyn FnMut(&[f32]),
|
||
) -> Result<(), DenoiseError> {
|
||
let n = self.tile;
|
||
let shape = ndarray::IxDyn(&[1, 1, n, n]);
|
||
// The vectors become the tensors: no copy on the way in.
|
||
let m = ort::value::Tensor::from_array(
|
||
ndarray::Array::from_shape_vec(shape.clone(), mosaic)
|
||
.map_err(|e| DenoiseError::Model(e.to_string()))?,
|
||
)?;
|
||
let s = ort::value::Tensor::from_array(
|
||
ndarray::Array::from_shape_vec(shape, sigma)
|
||
.map_err(|e| DenoiseError::Model(e.to_string()))?,
|
||
)?;
|
||
let acquired = self.model.acquire()?;
|
||
let mut session = acquired.lock();
|
||
let outputs = session.run(ort::inputs!["mosaic" => m, "sigma" => s])?;
|
||
let (shape, data) = outputs[0].try_extract_tensor::<f32>()?;
|
||
let dims: Vec<i64> = shape.iter().copied().collect();
|
||
if dims != [1, 3, n as i64, n as i64] {
|
||
return Err(DenoiseError::Model(format!(
|
||
"output is {dims:?}, expected [1, 3, {n}, {n}]"
|
||
)));
|
||
}
|
||
// And none on the way out: the frame is written from the runtime's buffer.
|
||
write(data);
|
||
Ok(())
|
||
}
|
||
}
|