How to Create a Mock ONNX Model in Python

Build and export a minimal mock ONNX model with a Reshape and Gemm layer using the onnx helper API.

Medium Python 3.9+ Aug 9, 2026 ML engineering pipelines 16 views 0 copies

onnx model-export mlops serialization deep-learning
Requires third-party packages — install first
pip install onnx numpy

Python code

60 lines
Python 3.9+ 
import onnx
import numpy as np
from onnx import helper, TensorProto

def create_mock_model():
    # Define input and output tensors
    input_tensor = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 224, 224])
    output_tensor = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 10])

    # Create a simple dummy layer (Gemm/FC)
    weight = helper.make_tensor(
        name='weight',
        data_type=TensorProto.FLOAT,
        dims=[10, 3*224*224],
        vals=np.random.randn(10, 3*224*224).astype(np.float32).flatten().tolist()
    )
    bias = helper.make_tensor(
        name='bias',
        data_type=TensorProto.FLOAT,
        dims=[10],
        vals=np.random.randn(10).astype(np.float32).tolist()
    )

    # Create nodes (reshape + gemm)
    reshape_node = helper.make_node(
        'Reshape',
        inputs=['input'],
        outputs=['flattened'],
        shape=[1, -1]
    )
    gemm_node = helper.make_node(
        'Gemm',
        inputs=['flattened', 'weight', 'bias'],
        outputs=['output'],
        alpha=1.0,
        beta=1.0,
        transB=1
    )

    # Build graph
    graph = helper.make_graph(
        [reshape_node, gemm_node],
        'mock_model',
        [input_tensor],
        [output_tensor]
    )

    # Create model
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 7  # IR version compatible with opset 13
    return model

if __name__ == "__main__":
    model = create_mock_model()
    onnx.save(model, "mock_model.onnx")
    print("Model created and exported")
    print(f"IR version: {model.ir_version}")
    print(f"Opset: {model.opset_import[0].version}")
    print(f"Number of nodes: {len(model.graph.node)}")
    print(f"Input name: {model.graph.input[0].name}, shape: {model.graph.input[0].type.tensor_type.shape.dim[0].dim_value}")

Output

stdout 
Model created and exported
IR version: 7
Opset: 13
Number of nodes: 2
Input name: input, shape: 1

How it works

The onnx.helper module provides constructors for tensors, nodes, and graphs, letting you define a model without training. make_graph assembles nodes with input/output value info, and make_model wraps it with opset metadata. The Reshape node flattens the 4D input to 2D, and Gemm applies a matrix multiplication with bias to produce 10 logits. Setting ir_version=7 matches opset 13 and ensures broad runtime compatibility.

Common mistakes

  • Forgetting to flatten weight dimensions to 1D for the Gemm node
  • Mismatching input tensor shape with the Reshape target dimension
  • Using an older IR version incompatible with the chosen opset

Variations

  1. Use `onnx.helper.make_tensor_value_info` with symbolic dimensions (e.g., 'N') for variable batch sizes
  2. Replace Gemm with Conv nodes for a CNN-style mock architecture

Real-world use cases

  • Testing ONNX Runtime loaders and inference engines with controlled dummy models before model deployment.
  • Verifying model validation, serialization, and versioning pipelines in CI/CD for ML artifacts.
  • Generating placeholder models for graph optimization, quantization, or conversion experiments without training real models.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from ML engineering pipelines

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes