Skip to main content

TensorModel

Trait TensorModel 

Source
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§

Source

fn name(&self) -> &str

The name clients use: /v2/models/{name}, /models/{name}/predict

Source

fn backend(&self) -> &str

The backend the model runs on, as reported by /model/info: cpu, wgpu, …

Source

fn inputs(&self) -> Vec<IoSpec>

The inputs, in the order TensorModel::run takes them

Source

fn outputs(&self) -> Vec<IoSpec>

The outputs, in the order TensorModel::run returns them

Source

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§

Source

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

Source

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".

Implementors§