add support for operations: reduceMean, add, mul, div, sigmoid
Some checks failed
Validate Operations / validate-operations (push) Failing after 51m52s

This commit is contained in:
NiccoloN
2026-03-30 15:41:12 +02:00
parent 5e7114f517
commit 39830be888
32 changed files with 1057 additions and 224 deletions

View File

@@ -1,5 +1,5 @@
#!/usr/bin/env python3
"""Generate ONNX test models for validating GEMM, Conv, Pooling, and Relu implementations."""
"""Generate ONNX test models for validating GEMM, Conv, Pooling, Relu, and ReduceMean implementations."""
import numpy as np
import onnx
@@ -19,102 +19,8 @@ def save_model(model, directory, filename):
print(f" {path.relative_to(OPERATIONS_DIR)}")
# ---------------------------------------------------------------------------
# GEMM tests
# ---------------------------------------------------------------------------
def gemm_non_square():
"""GEMM with non-square weight matrix: [B, K] @ [K, N], K != N."""
B, K, N = 4, 128, 64
W = numpy_helper.from_array(np.random.default_rng(42).uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_non_square", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/non_square", "gemm_non_square.onnx")
def gemm_with_bias():
"""GEMM with bias: Y = A @ W + C."""
B, K, N = 4, 128, 128
rng = np.random.default_rng(43)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"])
graph = helper.make_graph([node], "gemm_with_bias", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/with_bias", "gemm_with_bias.onnx")
def gemm_transB():
"""GEMM with transB=1: Y = A @ W^T."""
B, K, N = 4, 128, 64
rng = np.random.default_rng(44)
# W stored as [N, K], transposed during computation
W = numpy_helper.from_array(rng.uniform(-1, 1, (N, K)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"], transB=1)
graph = helper.make_graph([node], "gemm_transB", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/transB", "gemm_transB.onnx")
def gemm_alpha_beta():
"""GEMM with alpha and beta: Y = 0.5 * A @ W + 0.25 * C."""
B, K, N = 4, 64, 64
rng = np.random.default_rng(45)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"], alpha=0.5, beta=0.25)
graph = helper.make_graph([node], "gemm_alpha_beta", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/alpha_beta", "gemm_alpha_beta.onnx")
def gemm_small():
"""Small GEMM: [2, 8] @ [8, 4]."""
B, K, N = 2, 8, 4
rng = np.random.default_rng(46)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_small", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/small", "gemm_small.onnx")
def gemm_large():
"""Larger GEMM: [8, 256] @ [256, 128]."""
B, K, N = 8, 256, 128
rng = np.random.default_rng(47)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_large", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/large", "gemm_large.onnx")
def gemm_transB_with_bias():
"""GEMM with transB and bias: Y = A @ W^T + C."""
B, K, N = 4, 128, 64
rng = np.random.default_rng(48)
W = numpy_helper.from_array(rng.uniform(-1, 1, (N, K)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"], transB=1)
graph = helper.make_graph([node], "gemm_transB_with_bias", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/transB_with_bias", "gemm_transB_with_bias.onnx")
def make_int64_initializer(name, values):
return numpy_helper.from_array(np.asarray(values, dtype=np.int64), name=name)
# ---------------------------------------------------------------------------
@@ -248,6 +154,104 @@ def conv_large_spatial():
save_model(model, "conv/large_spatial", "conv_large_spatial.onnx")
# ---------------------------------------------------------------------------
# GEMM tests
# ---------------------------------------------------------------------------
def gemm_non_square():
"""GEMM with non-square weight matrix: [B, K] @ [K, N], K != N."""
B, K, N = 4, 128, 64
W = numpy_helper.from_array(np.random.default_rng(42).uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_non_square", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/non_square", "gemm_non_square.onnx")
def gemm_with_bias():
"""GEMM with bias: Y = A @ W + C."""
B, K, N = 4, 128, 128
rng = np.random.default_rng(43)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"])
graph = helper.make_graph([node], "gemm_with_bias", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/with_bias", "gemm_with_bias.onnx")
def gemm_transB():
"""GEMM with transB=1: Y = A @ W^T."""
B, K, N = 4, 128, 64
rng = np.random.default_rng(44)
# W stored as [N, K], transposed during computation
W = numpy_helper.from_array(rng.uniform(-1, 1, (N, K)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"], transB=1)
graph = helper.make_graph([node], "gemm_transB", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/transB", "gemm_transB.onnx")
def gemm_alpha_beta():
"""GEMM with alpha and beta: Y = 0.5 * A @ W + 0.25 * C."""
B, K, N = 4, 64, 64
rng = np.random.default_rng(45)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"], alpha=0.5, beta=0.25)
graph = helper.make_graph([node], "gemm_alpha_beta", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/alpha_beta", "gemm_alpha_beta.onnx")
def gemm_small():
"""Small GEMM: [2, 8] @ [8, 4]."""
B, K, N = 2, 8, 4
rng = np.random.default_rng(46)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_small", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/small", "gemm_small.onnx")
def gemm_large():
"""Larger GEMM: [8, 256] @ [256, 128]."""
B, K, N = 8, 256, 128
rng = np.random.default_rng(47)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W"], ["Y"])
graph = helper.make_graph([node], "gemm_large", [A], [Y], initializer=[W])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/large", "gemm_large.onnx")
def gemm_transB_with_bias():
"""GEMM with transB and bias: Y = A @ W^T + C."""
B, K, N = 4, 128, 64
rng = np.random.default_rng(48)
W = numpy_helper.from_array(rng.uniform(-1, 1, (N, K)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
node = helper.make_node("Gemm", ["A", "W", "C"], ["Y"], transB=1)
graph = helper.make_graph([node], "gemm_transB_with_bias", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "gemm/transB_with_bias", "gemm_transB_with_bias.onnx")
# ---------------------------------------------------------------------------
# Pooling tests
# ---------------------------------------------------------------------------
@@ -327,6 +331,55 @@ def maxpool_after_conv():
save_model(model, "pool/max_after_conv", "maxpool_after_conv.onnx")
# ---------------------------------------------------------------------------
# ReduceMean tests
# ---------------------------------------------------------------------------
def reducemean_basic():
"""ReduceMean over the feature dimension, preserving rank."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 1])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[1], keepdims=1)
graph = helper.make_graph([node], "reducemean_basic", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/basic", "reduce_mean_basic.onnx")
def reducemean_keepdims_0():
"""ReduceMean over the feature dimension, dropping the reduced axis."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[1], keepdims=0)
graph = helper.make_graph([node], "reducemean_keepdims_0", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/keepdims_0", "reduce_mean_keepdims_0.onnx")
def reducemean_4d_spatial():
"""ReduceMean over H and W on an NCHW tensor."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 4, 4])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 1, 1])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[2, 3], keepdims=1)
graph = helper.make_graph([node], "reducemean_4d_spatial", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/4d_spatial", "reduce_mean_4d_spatial.onnx")
def reducemean_after_conv():
"""Conv followed by ReduceMean over the spatial dimensions."""
rng = np.random.default_rng(62)
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 5, 5])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 1, 1])
W = numpy_helper.from_array(rng.uniform(-1, 1, (2, 3, 3, 3)).astype(np.float32), name="W")
B = numpy_helper.from_array(rng.uniform(-1, 1, (2,)).astype(np.float32), name="B")
conv = helper.make_node("Conv", ["X", "W", "B"], ["C"],
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
reduce = helper.make_node("ReduceMean", ["C"], ["Y"], axes=[2, 3], keepdims=1)
graph = helper.make_graph([conv, reduce], "reducemean_after_conv", [X], [Y], initializer=[W, B])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/after_conv", "reduce_mean_after_conv.onnx")
# ---------------------------------------------------------------------------
# Relu tests
# ---------------------------------------------------------------------------
@@ -381,6 +434,220 @@ def relu_after_gemm():
save_model(model, "relu/after_gemm", "relu_after_gemm.onnx")
# ---------------------------------------------------------------------------
# Sigmoid tests
# ---------------------------------------------------------------------------
def sigmoid_basic():
"""Standalone Sigmoid on a simple 2D tensor."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
node = helper.make_node("Sigmoid", ["X"], ["Y"])
graph = helper.make_graph([node], "sigmoid_basic", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "sigmoid/basic", "sigmoid_basic.onnx")
def sigmoid_4d():
"""Standalone Sigmoid on an NCHW tensor."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 4, 4])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 4, 4])
node = helper.make_node("Sigmoid", ["X"], ["Y"])
graph = helper.make_graph([node], "sigmoid_4d", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "sigmoid/4d", "sigmoid_4d.onnx")
def sigmoid_after_gemm():
"""Gemm followed by Sigmoid."""
B, K, N = 4, 64, 32
rng = np.random.default_rng(63)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
gemm = helper.make_node("Gemm", ["A", "W", "C"], ["G"])
sigmoid = helper.make_node("Sigmoid", ["G"], ["Y"])
graph = helper.make_graph([gemm, sigmoid], "sigmoid_after_gemm", [A], [Y], initializer=[W, C])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "sigmoid/after_gemm", "sigmoid_after_gemm.onnx")
# ---------------------------------------------------------------------------
# Add tests
# ---------------------------------------------------------------------------
def add_basic():
"""Elementwise Add on two inputs with identical shapes."""
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [4, 8])
B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
node = helper.make_node("Add", ["A", "B"], ["Y"])
graph = helper.make_graph([node], "add_basic", [A, B], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "add/basic", "add_basic.onnx")
def add_broadcast_row():
"""Elementwise Add with row-vector broadcasting."""
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
B = numpy_helper.from_array(np.random.default_rng(64).uniform(-1, 1, (8,)).astype(np.float32), name="B")
node = helper.make_node("Add", ["A", "B"], ["Y"])
graph = helper.make_graph([node], "add_broadcast_row", [A], [Y], initializer=[B])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "add/broadcast_row", "add_broadcast_row.onnx")
def add_after_gemm():
"""Gemm followed by Add with a broadcast bias vector."""
B, K, N = 4, 64, 32
rng = np.random.default_rng(65)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
D = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="D")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
gemm = helper.make_node("Gemm", ["A", "W", "C"], ["G"])
add = helper.make_node("Add", ["G", "D"], ["Y"])
graph = helper.make_graph([gemm, add], "add_after_gemm", [A], [Y], initializer=[W, C, D])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "add/after_gemm", "add_after_gemm.onnx")
# ---------------------------------------------------------------------------
# Mul tests
# ---------------------------------------------------------------------------
def mul_basic():
"""Elementwise Mul on two inputs with identical shapes."""
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [4, 8])
B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
node = helper.make_node("Mul", ["A", "B"], ["Y"])
graph = helper.make_graph([node], "mul_basic", [A, B], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "mul/basic", "mul_basic.onnx")
def mul_scalar_constant():
"""Elementwise Mul with scalar broadcasting."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
S = numpy_helper.from_array(np.asarray([1.5], dtype=np.float32), name="S")
node = helper.make_node("Mul", ["X", "S"], ["Y"])
graph = helper.make_graph([node], "mul_scalar_constant", [X], [Y], initializer=[S])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "mul/scalar_constant", "mul_scalar_constant.onnx")
def mul_after_conv():
"""Conv followed by Mul with per-channel scaling."""
rng = np.random.default_rng(66)
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 5, 5])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 3, 3])
W = numpy_helper.from_array(rng.uniform(-1, 1, (2, 3, 3, 3)).astype(np.float32), name="W")
B = numpy_helper.from_array(rng.uniform(-1, 1, (2,)).astype(np.float32), name="B")
S = numpy_helper.from_array(rng.uniform(0.5, 1.5, (1, 2, 1, 1)).astype(np.float32), name="S")
conv = helper.make_node("Conv", ["X", "W", "B"], ["C"],
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
mul = helper.make_node("Mul", ["C", "S"], ["Y"])
graph = helper.make_graph([conv, mul], "mul_after_conv", [X], [Y], initializer=[W, B, S])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "mul/after_conv", "mul_after_conv.onnx")
# ---------------------------------------------------------------------------
# Div tests
# ---------------------------------------------------------------------------
def div_basic():
"""Elementwise Div by a same-shape constant tensor."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
D = numpy_helper.from_array(np.random.default_rng(67).uniform(0.5, 2.0, (4, 8)).astype(np.float32), name="D")
node = helper.make_node("Div", ["X", "D"], ["Y"])
graph = helper.make_graph([node], "div_basic", [X], [Y], initializer=[D])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "div/basic", "div_basic.onnx")
def div_scalar_constant():
"""Elementwise Div with scalar broadcasting."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
S = numpy_helper.from_array(np.asarray([2.0], dtype=np.float32), name="S")
node = helper.make_node("Div", ["X", "S"], ["Y"])
graph = helper.make_graph([node], "div_scalar_constant", [X], [Y], initializer=[S])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "div/scalar_constant", "div_scalar_constant.onnx")
def div_after_gemm():
"""Gemm followed by Div with a broadcast divisor vector."""
B, K, N = 4, 64, 32
rng = np.random.default_rng(68)
W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
D = numpy_helper.from_array(rng.uniform(0.5, 2.0, (N,)).astype(np.float32), name="D")
A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
gemm = helper.make_node("Gemm", ["A", "W", "C"], ["G"])
div = helper.make_node("Div", ["G", "D"], ["Y"])
graph = helper.make_graph([gemm, div], "div_after_gemm", [A], [Y], initializer=[W, C, D])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "div/after_gemm", "div_after_gemm.onnx")
# ---------------------------------------------------------------------------
# ReduceMean tests
# ---------------------------------------------------------------------------
def reducemean_basic():
"""ReduceMean over the feature dimension, preserving rank."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 1])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[1], keepdims=1)
graph = helper.make_graph([node], "reducemean_basic", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/basic", "reduce_mean_basic.onnx")
def reducemean_keepdims_0():
"""ReduceMean over the feature dimension, dropping the reduced axis."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[1], keepdims=0)
graph = helper.make_graph([node], "reducemean_keepdims_0", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/keepdims_0", "reduce_mean_keepdims_0.onnx")
def reducemean_4d_spatial():
"""ReduceMean over H and W on an NCHW tensor."""
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 4, 4])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 1, 1])
node = helper.make_node("ReduceMean", ["X"], ["Y"], axes=[2, 3], keepdims=1)
graph = helper.make_graph([node], "reducemean_4d_spatial", [X], [Y])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/4d_spatial", "reduce_mean_4d_spatial.onnx")
def reducemean_after_conv():
"""Conv followed by ReduceMean over the spatial dimensions."""
rng = np.random.default_rng(62)
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 5, 5])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 1, 1])
W = numpy_helper.from_array(rng.uniform(-1, 1, (2, 3, 3, 3)).astype(np.float32), name="W")
B = numpy_helper.from_array(rng.uniform(-1, 1, (2,)).astype(np.float32), name="B")
conv = helper.make_node("Conv", ["X", "W", "B"], ["C"],
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
reduce = helper.make_node("ReduceMean", ["C"], ["Y"], axes=[2, 3], keepdims=1)
graph = helper.make_graph([conv, reduce], "reducemean_after_conv", [X], [Y], initializer=[W, B])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
save_model(model, "reduce_mean/after_conv", "reduce_mean_after_conv.onnx")
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
@@ -415,10 +682,36 @@ if __name__ == "__main__":
avgpool_include_pad()
maxpool_after_conv()
print("\nGenerating ReduceMean tests:")
reducemean_basic()
reducemean_keepdims_0()
reducemean_4d_spatial()
reducemean_after_conv()
print("\nGenerating Relu tests:")
relu_basic()
relu_4d()
relu_after_conv()
relu_after_gemm()
print("\nGenerating Sigmoid tests:")
sigmoid_basic()
sigmoid_4d()
sigmoid_after_gemm()
print("\nGenerating Add tests:")
add_basic()
add_broadcast_row()
add_after_gemm()
print("\nGenerating Mul tests:")
mul_basic()
mul_scalar_constant()
mul_after_conv()
print("\nGenerating Div tests:")
div_basic()
div_scalar_constant()
div_after_gemm()
print("\nDone.")