anira-project/stateful-lstm

Stateful LSTM for anira

A stateful LSTM guitar effect model (the architecture of Wright et al., "Real-Time Guitar Amplifier Emulation with Deep Learning", 2020), with its state as tensors, for ONNX Runtime, LiteRT, LibTorch, ExecuTorch (CPU, Core ML, Vulkan) and Core ML. anira's stateful network example.

Files

File	Engine	Size	SHA256
onnx/stateful_lstm.onnx	ONNX Runtime	10,501	f75e9d60cd3fc4ffa1282bcdd5a1fed46306647ba4eb3b51a61079b45ace81e8
litert/stateful_lstm.tflite	LiteRT, TensorFlow Lite	5,507,436	3eb200c608cee4aa5737775147b82ab02085f3cb7368103c8faf6f21613c248e
libtorch/stateful-lstm-dynamic.pt	LibTorch	21,769	277f262cab865697c102bdff91e840bb69b7f82cfeebea014577cdc0c6a74ab7
libtorch/stateful_lstm.pt	LibTorch	15,808	47ea81745158548b19938f5b4eafb9f235dae0469bf0bf25b8251d524e82a38f
executorch/stateful_lstm.pte	ExecuTorch 1.3.1, CPU (XNNPACK)	17,668	de3258a6f1825dc7dc14fde7e035f1eb26ebe966b5fc460ab10c92f21d6acdbe
executorch/stateful_lstm_coreml.pte	ExecuTorch 1.3.1 (Core ML, Apple)	34,807	3475f948e5f87eebbfe33607b622af83164d231f74049517abf85a88400971ae
executorch/stateful_lstm_vulkan.pte	ExecuTorch 1.3.1 (Vulkan, GPUs)	18,308	5432558a5162b4d6dbcdc43deab578bd92b0c98fcc7fde696609788bddc63ff1
coreml/stateful_lstm.mlpackage	Core ML (native)	14,045	per file in the catalog

Source

vackva/stateful-lstm @ 3ce7bec, models/model_0/; the training and export notebooks are in that repository.

State

The LSTM's state (h and c, 20 each) is a tensor pair: state_in [1,40] in, state_out [1,40] out, zero at the start of a stream, the previous call's state_out fed back as the next state_in. The original TorchScript model (libtorch/stateful-lstm-dynamic.pt) keeps the state inside (self.hidden), which only LibTorch carries between calls: its ONNX, LiteRT and ExecuTorch exports started every block from a zero state. These, now in upstream/, are replaced by exports of export/lstm.py, the same LSTM with its state as tensors, which match the original over consecutive blocks exactly (ONNX 5.7e-7). The original stays for LibTorch (stateful-lstm-libtorch in anira foundation models).

Inputs	Outputs
audio_in [2048,1,1] (time, batch, channel), state_in [1,40]	audio_out [2048,1,1], state_out [1,40]

Formats and checks

Every format is checked against ONNX Runtime over twelve consecutive blocks of guitar-like audio with the state fed back (export/check.py), and anira foundation models streams every format over twelve consecutive calls (tools/check_streaming.py).

Format	Files	Max deviation from ONNX Runtime
ONNX Runtime (against the original TorchScript, 20 blocks)	onnx/stateful_lstm.onnx	5.7e-7
LibTorch (traced, explicit state)	libtorch/stateful_lstm.pt	0 (against the original)
LiteRT, TensorFlow Lite	litert/stateful_lstm.tflite	7.3e-7
ExecuTorch, CPU (XNNPACK and a scan loop, 2.9 ms per 2048-sample call)	executorch/stateful_lstm.pte	6.7e-7
ExecuTorch, Core ML (Apple, float32, 0.85 ms per call)	executorch/stateful_lstm_coreml.pte	5.2e-7
ExecuTorch, Vulkan (GPUs)	executorch/stateful_lstm_vulkan.pte	exported; not run here (no Vulkan runtime)
Core ML ML program (anira's native engine, coming with v3)	coreml/stateful_lstm.mlpackage	5.2e-7

torch.export unrolls torch.nn.LSTM over the block's 2048 steps: thousands of small operators, which ExecuTorch's delegates lower in gigabytes and run slowly (Core ML took 190 ms per call). The ExecuTorch programs avoid that. The CPU and Vulkan programs take ScanLstm (export/lstm.py), the same LSTM with its recurrence as torch.scan, which ExecuTorch runs as a loop over one step: the delegate takes the input projection and the output layer, the loop runs on ExecuTorch's CPU kernels (a 20-unit step is too small for a GPU). The Core ML program keeps the LSTM as one operator that converts to Core ML's lstm, as the native Core ML package does. LiteRT unrolls the LSTM too and takes 1.9 ms per call. The CPU program also runs on the ExecuTorch runtime anira ships (anira-project/backends v2.4.0) and matches PyTorch there over consecutive calls.

cd export
uv run --no-project export.py <stateful-lstm-dynamic.pt> <prefix>              # ONNX, LibTorch
uv run --no-project export_litert.py <stateful-lstm-dynamic.pt> <prefix>
uv run --no-project export_executorch.py <stateful-lstm-dynamic.pt> <prefix> [--coreml | --vulkan]
uv run --no-project export_coreml.py <stateful-lstm-dynamic.pt> <prefix>

upstream/ keeps anira's original ONNX, LiteRT and ExecuTorch files (upstream/), which the descriptions no longer name: they lose the state between calls (above).

This repository holds the weights, the export scripts and the licence. The anira model descriptions live with the code in anira foundation models (models/stateful-lstm), versioned with git. Pin a revision and check each file's SHA256 when you download:

https://huggingface.co/anira-project/stateful-lstm/resolve/<revision>/onnx/stateful_lstm.onnx

Licence

GPL-3.0: LICENSE. The model is trained on GuitarML's TS9 recording, which is GPL-3.0; see NOTICE.md. Software that ships these weights may have to meet the GPL's terms; check before you do.