Skip to content

Dynamic shapes

An engine normally accepts exactly the input shapes it was built with. Trace a lidar model on a 100-voxel frame and the engine takes 100 voxels — forever. Real frames aren't that obliging: point counts change every sweep, batch sizes change per request, and padding every input up to a worst case wastes latency on every frame that isn't the worst case.

Dynamic shapes let you nominate an axis that varies. You give it a range, and the engine accepts any size inside it — one build, every frame size, no padding.

graph LR
    A[1. Pick the varying axes] --> B[2. prepare<br/>--dynamic-dim] --> C[3. compile] --> D[4. Feed any size<br/>in range]

You opt in at prepare time with --dynamic-dim. Everything after that is unchanged: compile and verify need no extra flags, because the range travels inside the ONNX.

Step 1 — Decide which axis varies

Look at your get_fwd_args() and ask which dimension changes frame to frame. For the sparse model in Compile your first model, it's the voxel count N — dimension 0 of both features and coords:

def get_fwd_args():
    ...
    features = torch.randn(n, 4)  # (N, 4)  <- N varies
    coords = torch.cat([torch.zeros(n, 1, dtype=torch.int32), zyx], 1)  # (N, 4)  <- same N
    return (features, coords)

Then pick the range. min and max are the sizes the engine will accept — set them to the real bounds of your data, with a little headroom, not to the size of your example.

Step 2 — Prepare with --dynamic-dim

Each --dynamic-dim marks one axis of one input, and the flag is repeatable:

--dynamic-dim '<input-index>:<dim>=<min>[..<opt>]..<max>[@<symbol>]'

<input-index> counts positions in the tuple get_fwd_args() returns — features is 0, coords is 1. <dim> then counts axes within that one input's own shape: features is (N, 4), so dimension 0 is the voxel count that varies and dimension 1 is the four features per voxel, which doesn't. So the sparse model, whose voxel count can run from a nearly empty frame to 20 000 voxels:

yasp-toolkit embedded prepare model.py \
  --onnx-out      build/model.onnx \
  --npz-out       build/model.npz \
  --reference-out build/model.reference.npz \
  --dynamic-dim   '0:0=1..20000@N' \
  --dynamic-dim   '1:0=1..20000@N'

Inputs you don't name keep the fixed size of their example, so a model with one varying axis needs one flag and nothing else changes.

Part Meaning
0:0 Input 0, dimension 0 — the axis being freed.
=1..20000 The engine accepts any size in [1, 20000].
=1..4000..20000 Same range, but tuned for 4000 — see choosing the tuned size.
@N The axis's name. Optional; see tying two inputs together.

Tying two inputs together

features and coords describe the same voxels, so their row counts are always equal. Giving both axes the symbol @N states that: the builder can use the equality, and the two dimensions appear in the ONNX under one shared name.

That declaration is a promise, not a check — nothing verifies it at runtime. Feed 4 000 rows of features and 4 001 rows of coords and each passes its own range check independently, with undefined results downstream. If two axes aren't genuinely locked together, give them different symbols (or omit @ entirely, which names each axis uniquely) and let them vary on their own.

Choosing the tuned size

The range says what the engine accepts. A second, separate thing says what it's tuned for: TensorRT picks kernels for one size, so the engine runs fastest there and merely correctly elsewhere. Get this wrong and you have a build that's technically fine and needlessly slow.

By default the tuned size is your example's. Two axes, one example frame of 100 voxels — that's what gets tuned for, and it's the right answer whenever your example is representative:

--dynamic-dim '0:0=1..20000@N'      # tuned for the example frame

Add a middle value to tune for something else. Useful when a convenient example frame isn't a typical one — a tiny smoke-test input, or a worst-case frame you keep around for bounds:

--dynamic-dim '0:0=1..4000..20000@N'   # accepts 1..20000, tuned for 4000

Read it as min, opt, max, low to high. All three must be ordered, so 1..25000..20000 is rejected rather than quietly retuning.

Whichever form you use:

  • The example's size must still fall inside the range. A frame of 100 voxels with --dynamic-dim '0:0=500..20000@N' is rejected at prepare, naming the input and the range — the example is still the frame that gets traced and referenced, so it has to be a shape the engine can run.
  • You can pin a tuned size your example isn't, and nothing objects. That's the point of the middle value.
  • Mix freely per axis. One --dynamic-dim can pin an opt while another leaves it to the example.

verify measures at the example's size, not the tuned size

Every number verify reports comes from a run on your example inputs: the accuracy scores, the benchmark table, and the per-layer profiling table alike. The worker captures its --benchmark and --profile-layers runs against the same sample inputs, and the torch host baseline row is timed on them too. Pinning a tuned size doesn't move any of that.

For accuracy that's usually fine — which kernel TensorRT picked doesn't change what the math computes.

For speed it can mislead badly, and speed is normally the whole reason to pin. Tune for 4000 while your example is 100 voxels and the latency you read back was measured at 100, on an engine optimized for 4000 — the one size where its kernel choice is least suited. Those milliseconds are not the milliseconds you'll see in deployment, and the comparison you'd naturally draw (against a build without the pin) is measuring the wrong thing.

If you want the reported latency to describe the size you tuned for, make an example frame that size and drop the middle value — then the default lines all three up: tuned, scored, and timed at the same shape.

Calibrating with several frames

When get_fwd_args() returns a list of frames, the first one is the frame the engine is tuned for (unless you pin a middle value, which overrides it). The rest still matter — on sparse models they set how far each layer's output can shrink — but they don't move the tuning point. Put your most representative frame first; nothing warns you if you don't.

Step 3 — Compile and verify, unchanged

The range is stamped into the ONNX, so the worker reads it straight off the model you upload:

yasp-toolkit embedded compile \
  --arch-spec "$ARCH" \
  --precision fp16 \
  --onnx      build/model.onnx \
  --npz       build/model.npz \
  --output    "build/model.$ARCH.fp16.tgz"

yasp-toolkit embedded verify \
  build/model.reference.npz \
  "build/model.$ARCH.fp16.tgz"

verify scores and times the engine at your example's size, exactly as for a static build — the reference outputs and the worker's captured benchmark both come from that shape, whether or not you pinned a different tuned size. See the warning above if those differ.

Step 4 — Feed a varying size at runtime

The engine now carries a range instead of a fixed shape. Set the actual shape before each inference, then run as usual:

// features and coords for a 4 217-voxel frame
binary.setInputShape(0, yasp::inference::Shape::make(4217, 4));
binary.setInputShape(1, yasp::inference::Shape::make(4217, 4));

A shape outside the built range is refused with ERROR_INVALID_SHAPE rather than producing garbage. Output buffers are sized from what the engine actually produces for that input, so an output whose size follows the input needs no declaration from you — it comes out at the right size on its own.

Sparse models get one extra effect

On a model with spconv layers, freeing an axis also changes how the per-layer capacities are chosen. Normally they're measured by running your example frames and recording how many voxels each layer produced. That measurement is exactly what a varying voxel count invalidates, so with --dynamic-dim the capacities are derived from the range instead, and each layer's output is bounded relative to its own input at runtime.

Two practical consequences:

  • --nnz-headroom no longer applies. It's a margin on a measurement that no longer happens.
  • Your calibration frames still matter. They set how far each sparse layer is allowed to shrink its output. Frames that aren't representative of deployment can bound a layer too tightly, so calibrate on frames that look like real data.

Not every model supports dynamic shapes yet

Some model families have no dynamic path and are rejected at prepare with a clear message rather than compiled into something that doesn't work. If you hit that, compile the model statically — pick the input size you actually deploy at — and tell support@yasp.ai which model you needed it for.

How it works

  • prepare --dynamic-dim marks those axes as free in the exported ONNX and stamps the extents into the model file as metadata, so nothing needs to travel beside it. An axis you didn't pin a tuned size for is stamped as just its range.
  • compile reads that metadata on the worker and turns it into TensorRT's optimization profile: your ranges as the accepted bounds, and for the tuned size either the value you pinned or — for an axis stamped with only its range — the shape of that input in the sample .npz. An ONNX with no such metadata compiles exactly as before.
  • The engine derives every output size from its inputs, so only inputs ever need a declared range.
  • At runtime the shipped binary checks each shape you set against the profile and allocates output buffers from the sizes the engine reports back.