Skip to content

ONNX Export and Inference

ONNX export and inference for trained Learners.

OnnxInferenceWrapper 

OnnxInferenceWrapper(path: str | Path, **session_options)

Run an exported ONNX model with onnxruntime. Same API as InferenceWrapper.

Parameters:

Name Type Description Default
path str | Path

path to the exported .onnx model file

 required
session_options

additional keyword arguments forwarded to onnxruntime.InferenceSession

 {}
Source code in tsfast/inference/onnx.py 
77
78
79
80
81
82
def __init__(self, path: str | Path, **session_options):
    import onnxruntime as ort

    self.session = ort.InferenceSession(str(path), **session_options)
    self._input_name = self.session.get_inputs()[0].name
    self._output_name = self.session.get_outputs()[0].name

inference 

inference(np_input: ndarray, np_output_init: ndarray | None = None) -> np.ndarray

Run inference on numpy input, returns numpy output. Output ndim mirrors input ndim.

Source code in tsfast/inference/onnx.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def inference(
    self,
    np_input: np.ndarray,  # input time series
    np_output_init: np.ndarray | None = None,  # initial output (for PredictionCallback models)
) -> np.ndarray:
    "Run inference on numpy input, returns numpy output. Output ndim mirrors input ndim."
    input_ndim = np_input.ndim
    u = self._prepare(np_input, "np_input")
    if np_output_init is not None:
        y_init = self._prepare(np_output_init, "np_output_init")
        if u.shape[0] != y_init.shape[0]:
            raise ValueError(
                f"Batch size mismatch: np_input has {u.shape[0]}, np_output_init has {y_init.shape[0]}."
            )
        seq_len = u.shape[1]
        if y_init.shape[1] != seq_len:
            y_init = (
                y_init[:, :seq_len, :]
                if y_init.shape[1] > seq_len
                else np.pad(y_init, ((0, 0), (0, seq_len - y_init.shape[1]), (0, 0)))
            )
        u = np.concatenate((u, y_init), axis=-1)

    result = self.session.run([self._output_name], {self._input_name: u})[0]
    match input_ndim:
        case 1:
            if result.shape[-1] != 1:
                raise ValueError(
                    f"Cannot return 1D output: model produces {result.shape[-1]} features. "
                    f"Pass 2D input (seq_len, features) instead."
                )
            return result[0, :, 0]
        case 2:
            return result[0]
        case _:
            return result

export_onnx 

export_onnx(learner, path: str | Path, opset_version: int = 17, seq_len: int | None = None) -> Path

Export a trained Learner's model to ONNX format with normalization baked in.

Source code in tsfast/inference/onnx.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def export_onnx(
    learner,  # trained Learner with .model and .dls
    path: str | Path,  # output .onnx file path
    opset_version: int = 17,  # ONNX opset version
    seq_len: int | None = None,  # override sequence length for dummy input (default: from dls)
) -> Path:
    "Export a trained Learner's model to ONNX format with normalization baked in."
    import onnx

    path = Path(path)
    if path.suffix != ".onnx":
        path = path.with_suffix(".onnx")

    model = learner.model
    _check_no_ar(model)

    model = model.cpu().eval()
    dummy = _get_dummy_input(learner, seq_len)

    torch.onnx.export(
        model,
        (dummy,),
        f=str(path),
        opset_version=opset_version,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch", 1: "seq_len"},
            "output": {0: "batch", 1: "seq_len"},
        },
        dynamo=False,
    )

    onnx_model = onnx.load(str(path))
    onnx.checker.check_model(onnx_model)
    return path