Export Formats

How KynML saves trained models: torch state dict, onnx cross-platform graph, and torchscript compiled script.

Overview

The export block is optional. When present it runs after training and evaluation:

export:
    format = <torch|onnx|torchscript>
    path = "models/output.pt"
    input_shape = [1, 10]   # required for onnx only
    opset = 17              # onnx only, default 17

Export happens in export_model() in the generated script. If no export block is present, export_model() is a no-op.

torch — state dict

The simplest format. Saves the model's parameter tensors as a Python dict using torch.save(model.state_dict(), path).

export:
    format = torch
    path = "models/house_price_model.pt"

Generated code

if EXPORT_FORMAT == "torch":
    torch.save(model.state_dict(), path)
    print(f"Saved model to {path}")

Loading

import torch
import torch.nn as nn

class HousePriceModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(10, 64), nn.ReLU(),
            nn.Linear(64, 32), nn.ReLU(),
            nn.Linear(32, 1),
        )
    def forward(self, x):
        return self.net(x)

model = HousePriceModel()
model.load_state_dict(torch.load("models/house_price_model.pt", map_location="cpu", weights_only=True))
model.eval()

Or use the KynML serving module to get a ready-made FastAPI app:

from kynml.parser import parse_file
from kynml.semantic import validate_program
from kynml.serving.generator import generate_service

program = parse_file("house_price.kyn")
validate_program(program)
generate_service(program, model_path="models/house_price_model.pt", out_dir="service/")

When to use

• Default choice. Compatible with everything in the PyTorch ecosystem.
• Requires the model class definition at load time — the generated .py file or the spec is your source of truth.
• Not portable to non-PyTorch runtimes (use ONNX or TorchScript for that).

onnx — cross-platform inference graph

Exports a static computation graph to the ONNX interchange format. No Python or PyTorch required at inference time — load with ONNX Runtime, TensorRT, CoreML, or any ONNX-compatible runtime.

export:
    format = onnx
    path = "models/churn_model.onnx"
    input_shape = [1, 15]
    opset = 17

input_shape is required for ONNX. The semantic validator rejects ONNX exports without it:

KynMLSemanticError: Export format 'onnx' requires input_shape

Generated code

elif EXPORT_FORMAT == "onnx":
    if EXPORT_INPUT_SHAPE is None:
        raise RuntimeError("ONNX export requires input_shape")
    dummy = torch.randn(*EXPORT_INPUT_SHAPE)
    torch.onnx.export(
        model,
        dummy,
        str(path),
        opset_version=EXPORT_OPSET,     # default 17
        input_names=["input"],
        output_names=["output"],
    )
    print(f"Saved ONNX model to {path}")

input_shape

A list of integers matching [batch_size, num_features]. The dummy tensor (torch.randn(*input_shape)) is traced through the model to build the ONNX graph. Typical shapes:

The batch dimension can be symbolic. To export with a dynamic batch axis, edit the generated script and pass dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}} to torch.onnx.export.

opset

ONNX operator set version. Default: 17. Minimum supported by most runtimes: 11. ONNX Runtime 1.16+ supports opset 18. Use the newest opset your target runtime supports.

Loading with ONNX Runtime

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("models/churn_model.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name

features = np.array([[1.0, 35.0, 2.0, 80.0, 1.0, 0.0, 12.0, 3.0, 1.0, 4.0, 0.0, 2.0, 7.0, 1.0, 0.0]], dtype=np.float32)
result = sess.run(None, {input_name: features})
print(result[0])  # [[0.1234]]

Install ONNX Runtime:

pip install onnxruntime          # CPU
pip install onnxruntime-gpu      # CUDA

When to use

• Deploying to a non-Python environment (mobile, embedded, browser via WASM).
• Serving with TensorRT for maximum GPU inference throughput.
• A/B testing different frameworks without retraining.
• Regulatory/audit contexts where you need a frozen, inspectable graph.

torchscript — compiled PyTorch script

Serialises the model via torch.jit.script, which traces the Python model class into a language-agnostic IR. The output is a self-contained .pt file that includes both weights and computation graph — no model class definition required at load time.

export:
    format = torchscript
    path = "models/deepnet.pt"

Generated code

elif EXPORT_FORMAT == "torchscript":
    scripted = torch.jit.script(model)
    scripted.save(str(path))
    print(f"Saved TorchScript model to {path}")

Note: torch.jit.script (not torch.jit.trace). Scripting handles control flow; tracing does not. All KynML-generated models use nn.Sequential with standard PyTorch modules, which script cleanly.

Loading

import torch

model = torch.jit.load("models/deepnet.pt", map_location="cpu")
model.eval()

x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])  # shape [1, input_size]
with torch.no_grad():
    prediction = model(x)
print(prediction)

No Python class definition needed. The file is portable between machines with the same PyTorch version.

Loading from C++

#include <torch/script.h>

auto model = torch::jit::load("models/deepnet.pt");
model.eval();

auto input = torch::ones({1, 10});
auto output = model.forward({input}).toTensor();
std::cout << output << std::endl;

When to use

• Deploying within a Python or C++ application where ONNX is not needed.
• Self-contained model files without needing the class source code at load time.
• Mobile deployment via torch.jit mobile interpreter (requires additional compilation step).

Format comparison

Full spec examples

Torch export (regression)

export:
    format = torch
    path = "models/house_price_model.pt"

ONNX export (binary classifier, 15 features)

export:
    format = onnx
    path = "models/churn_model.onnx"
    input_shape = [1, 15]
    opset = 17

TorchScript export (multiclass, 4 features)

export:
    format = torchscript
    path = "models/iris_classifier.pt"

ONNX with older opset for maximum runtime compatibility

export:
    format = onnx
    path = "models/legacy_model.onnx"
    input_shape = [1, 10]
    opset = 12

Troubleshooting

Export format 'onnx' requires input_shape — add input_shape = [batch, features] to your export block. The batch dimension is typically 1 for the dummy input.

torch.jit.frontend.UnsupportedNodeError — a module in your model is not scriptable. All standard KynML layers (Dense, Dropout, BatchNorm1d, standard activations) are scriptable. Custom post-compile edits using Lambda layers or dynamic Python are not.

opset version X not supported — your installed ONNX Runtime version does not support the opset you specified. Lower opset (e.g. from 17 to 12) or upgrade onnxruntime.

Model exported but predictions differ — check that normalize = true was applied consistently. The StandardScaler fit during training is not saved in the exported model file. For inference, normalise inputs manually using the same mean/std before passing to the model.