svod_model/resnet/jit.rs
1//! JIT wrapper for [`ResNet`]. Compiles the full forward graph once at
2//! `prepare()` time and replays it per call with the batch dimension bound to
3//! the live value of `b` through
4//! [`execute_with_vars`](ResNetJit::execute_with_vars).
5//!
6//! Shape contract:
7//! - `prepare(InputSpec::f32(&[max_b, 3, H, W]))` bakes `H`/`W` into the plan.
8//! - `b` is rebindable on every call; values are clamped to
9//! `[1, max_batch_size]` by the macro-generated setters.
10//! - For a different image resolution, prepare a fresh wrapper or call
11//! `prepare` again on this one with a new `InputSpec`.
12//!
13//! See `website/docs/architecture/jit-graphs.md` for the wrapper contract.
14
15extern crate self as svod_model;
16
17use svod_macros::jit_wrapper;
18
19use super::model::ResNet;
20
21jit_wrapper! {
22 ResNetJit(ResNet) {
23 images: Tensor,
24
25 vars {
26 b: (1, model.config.max_batch_size),
27 }
28
29 build(images, b) {
30 model.forward(images, &b)
31 }
32 }
33}