pub trait TensorModel:
Send
+ Sync
+ Debug {
// Required methods
fn name(&self) -> &str;
fn backend(&self) -> &str;
fn inputs(&self) -> Vec<IoSpec>;
fn outputs(&self) -> Vec<IoSpec>;
fn run(&self, inputs: Vec<HostTensor>) -> Result<Vec<HostTensor>>;
// Provided methods
fn fixed_batch(&self) -> Option<usize> { ... }
fn int_range(&self) -> (i64, i64) { ... }
}Expand description
A model the server can run: what it takes, what it returns, and how to run it.
The one trait between the server and the engine that runs a model. furnace-burn implements
it for Burn models (furnace_burn::custom, furnace_burn::from_fn); implement it yourself
to serve a model from any other engine, with every feature of the server (named tensors over
the v2 protocol, batching, metrics, reloads, …), through Model::custom.
use furnace::{DataType, HostTensor, IoSpec, Model, TensorModel, Values};
/// Doubles every value
#[derive(Debug)]
struct Doubler;
impl TensorModel for Doubler {
fn name(&self) -> &str { "doubler" }
fn backend(&self) -> &str { "cpu" }
fn inputs(&self) -> Vec<IoSpec> { vec![IoSpec::fp32("input", &[4])] }
fn outputs(&self) -> Vec<IoSpec> { vec![IoSpec::fp32("output", &[4])] }
fn run(&self, inputs: Vec<HostTensor>) -> furnace::Result<Vec<HostTensor>> {
// Checked already: one FP32 tensor of [batch, 4]
let HostTensor { shape, values: Values::F32(values) } = &inputs[0] else {
unreachable!("checked against inputs()")
};
let doubled = values.iter().map(|v| v * 2.0).collect();
Ok(vec![HostTensor { shape: shape.clone(), values: Values::F32(doubled) }])
}
}
let model = Model::custom(Doubler);
assert_eq!(model.predict(vec![1.0, 2.0, 3.0, 4.0]).unwrap(), [2.0, 4.0, 6.0, 8.0]);Required Methods§
Sourcefn backend(&self) -> &str
fn backend(&self) -> &str
The backend the model runs on, as reported by /model/info: cpu, wgpu, …
Sourcefn inputs(&self) -> Vec<IoSpec>
fn inputs(&self) -> Vec<IoSpec>
The inputs, in the order TensorModel::run takes them
Sourcefn outputs(&self) -> Vec<IoSpec>
fn outputs(&self) -> Vec<IoSpec>
The outputs, in the order TensorModel::run returns them
Sourcefn run(&self, inputs: Vec<HostTensor>) -> Result<Vec<HostTensor>>
fn run(&self, inputs: Vec<HostTensor>) -> Result<Vec<HostTensor>>
Run one batch. The server has already checked the inputs against
TensorModel::inputs: one tensor each, in that order, with the batch as the first
dimension, every fixed dimension as declared, and the same batch size throughout.
Blocking; the server calls it off the async runtime.
Provided Methods§
Sourcefn fixed_batch(&self) -> Option<usize>
fn fixed_batch(&self) -> Option<usize>
The batch size the model was built for, if it takes no other: larger batches are then run in chunks of it
Sourcefn int_range(&self) -> (i64, i64)
fn int_range(&self) -> (i64, i64)
The integers the engine can hold. Integer inputs outside this range are refused before
TensorModel::run sees them.
Dyn Compatibility§
This trait is dyn compatible.
In older versions of Rust, dyn compatibility was called "object safety".