718 lines
36 KiB
Python
718 lines
36 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate ONNX test models for validating GEMM, Conv, Pooling, Relu, and ReduceMean implementations."""
|
|
|
|
import numpy as np
|
|
import onnx
|
|
from onnx import helper, TensorProto, numpy_helper
|
|
from pathlib import Path
|
|
|
|
OPERATIONS_DIR = Path(__file__).parent
|
|
|
|
|
|
def save_model(model, directory, filename):
|
|
"""Save an ONNX model, creating the directory if needed."""
|
|
d = OPERATIONS_DIR / directory
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
path = d / filename
|
|
onnx.checker.check_model(model)
|
|
onnx.save(model, str(path))
|
|
print(f" {path.relative_to(OPERATIONS_DIR)}")
|
|
|
|
|
|
def make_int64_initializer(name, values):
|
|
return numpy_helper.from_array(np.asarray(values, dtype=np.int64), name=name)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conv tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def conv_3x3_kernel():
|
|
"""Conv with 3x3 kernel, no padding."""
|
|
# Input: [1, 1, 5, 5], Kernel: [1, 1, 3, 3] -> Output: [1, 1, 3, 3]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5, 5])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 3, 3])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(50).uniform(-1, 1, (1, 1, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_3x3", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/kernel_3x3", "conv_kernel_3x3.onnx")
|
|
|
|
|
|
def conv_stride2():
|
|
"""Conv with 3x3 kernel and stride 2."""
|
|
# Input: [1, 1, 6, 6], Kernel: [1, 1, 3, 3], stride 2 -> Output: [1, 1, 2, 2]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 6, 6])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 2, 2])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(51).uniform(-1, 1, (1, 1, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[2, 2], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_stride2", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/stride_2", "conv_stride_2.onnx")
|
|
|
|
|
|
def conv_multi_channel():
|
|
"""Conv with multiple input and output channels."""
|
|
# Input: [1, 3, 5, 5], Kernel: [4, 3, 3, 3] -> Output: [1, 4, 3, 3]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 5, 5])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 3, 3])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(52).uniform(-1, 1, (4, 3, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_multi_channel", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/multi_channel", "conv_multi_channel.onnx")
|
|
|
|
|
|
def conv_1x1():
|
|
"""1x1 pointwise convolution (channel mixing)."""
|
|
# Input: [1, 8, 4, 4], Kernel: [4, 8, 1, 1] -> Output: [1, 4, 4, 4]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 8, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 4, 4])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(53).uniform(-1, 1, (4, 8, 1, 1)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_1x1", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/pointwise_1x1", "conv_1x1.onnx")
|
|
|
|
|
|
def conv_same_padding_3x3():
|
|
"""Conv 3x3 with SAME_UPPER padding, preserving spatial dimensions."""
|
|
# Input: [1, 1, 5, 5], Kernel: [1, 1, 3, 3], SAME_UPPER -> Output: [1, 1, 5, 5]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5, 5])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 5, 5])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(54).uniform(-1, 1, (1, 1, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], auto_pad="SAME_UPPER")
|
|
graph = helper.make_graph([node], "conv_same_3x3", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/same_padding_3x3", "conv_same_padding_3x3.onnx")
|
|
|
|
|
|
def conv_explicit_padding():
|
|
"""Conv 3x3 with explicit asymmetric padding."""
|
|
# Input: [1, 1, 4, 4], Kernel: [1, 1, 3, 3], pads=[1,1,1,1] -> Output: [1, 1, 4, 4]
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 4])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(55).uniform(-1, 1, (1, 1, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1])
|
|
graph = helper.make_graph([node], "conv_explicit_pad", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/explicit_padding", "conv_explicit_padding.onnx")
|
|
|
|
|
|
def conv_with_bias_3x3():
|
|
"""Conv 3x3 with bias."""
|
|
# Input: [1, 3, 5, 5], Kernel: [2, 3, 3, 3], Bias: [2] -> Output: [1, 2, 3, 3]
|
|
rng = np.random.default_rng(56)
|
|
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")
|
|
node = helper.make_node("Conv", ["X", "W", "B"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_with_bias_3x3", [X], [Y], initializer=[W, B])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/with_bias_3x3", "conv_with_bias_3x3.onnx")
|
|
|
|
|
|
def conv_batch_2():
|
|
"""Batched conv (batch=2) with SAME_UPPER padding and bias."""
|
|
# Input: [2, 3, 3, 3], Kernel: [1, 3, 2, 2], Bias: [1] -> Output: [2, 1, 3, 3]
|
|
rng = np.random.default_rng(57)
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 3, 3])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 1, 3, 3])
|
|
W = numpy_helper.from_array(rng.uniform(-1, 1, (1, 3, 2, 2)).astype(np.float32), name="W")
|
|
B = numpy_helper.from_array(rng.uniform(-1, 1, (1,)).astype(np.float32), name="B")
|
|
node = helper.make_node("Conv", ["X", "W", "B"], ["Y"],
|
|
kernel_shape=[2, 2], strides=[1, 1], auto_pad="SAME_UPPER")
|
|
graph = helper.make_graph([node], "conv_batch_2", [X], [Y], initializer=[W, B])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "conv/batch_2", "conv_batch_2.onnx")
|
|
|
|
|
|
def conv_large_spatial():
|
|
"""Conv on larger spatial input: [1, 1, 8, 8] with 3x3 kernel."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 8, 8])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 6, 6])
|
|
W = numpy_helper.from_array(
|
|
np.random.default_rng(58).uniform(-1, 1, (1, 1, 3, 3)).astype(np.float32), name="W")
|
|
node = helper.make_node("Conv", ["X", "W"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "conv_large_spatial", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def maxpool_basic():
|
|
"""MaxPool 2x2 with stride 1."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 3, 3])
|
|
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "maxpool_basic", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/max_basic", "maxpool_basic.onnx")
|
|
|
|
|
|
def maxpool_stride2_multichannel():
|
|
"""MaxPool 2x2 with stride 2 on multiple channels."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 5, 6, 6])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 5, 3, 3])
|
|
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "maxpool_stride2_multichannel", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/max_stride2_multichannel", "maxpool_stride2_multichannel.onnx")
|
|
|
|
|
|
def maxpool_same_upper():
|
|
"""MaxPool 3x3 with SAME_UPPER padding."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 5, 5])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 3, 3])
|
|
node = helper.make_node("MaxPool", ["X"], ["Y"], kernel_shape=[3, 3], strides=[2, 2], auto_pad="SAME_UPPER")
|
|
graph = helper.make_graph([node], "maxpool_same_upper", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/max_same_upper", "maxpool_same_upper.onnx")
|
|
|
|
|
|
def avgpool_basic():
|
|
"""AveragePool 2x2 with stride 1."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 3, 3])
|
|
node = helper.make_node("AveragePool", ["X"], ["Y"], kernel_shape=[2, 2], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([node], "avgpool_basic", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/avg_basic", "avgpool_basic.onnx")
|
|
|
|
|
|
def avgpool_explicit_padding():
|
|
"""AveragePool 3x3 with explicit padding, excluding pad from the divisor."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2])
|
|
node = helper.make_node("AveragePool", ["X"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=0)
|
|
graph = helper.make_graph([node], "avgpool_explicit_padding", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/avg_explicit_padding", "avgpool_explicit_padding.onnx")
|
|
|
|
|
|
def avgpool_include_pad():
|
|
"""AveragePool 3x3 with explicit padding, including pad in the divisor."""
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 2, 4, 4])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 2, 2])
|
|
node = helper.make_node("AveragePool", ["X"], ["Y"],
|
|
kernel_shape=[3, 3], strides=[2, 2], pads=[1, 1, 1, 1], count_include_pad=1)
|
|
graph = helper.make_graph([node], "avgpool_include_pad", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "pool/avg_include_pad", "avgpool_include_pad.onnx")
|
|
|
|
|
|
def maxpool_after_conv():
|
|
"""Conv followed by MaxPool to validate pooling on lowered conv results."""
|
|
rng = np.random.default_rng(59)
|
|
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 6, 6])
|
|
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 2, 2])
|
|
W = numpy_helper.from_array(rng.uniform(-1, 1, (4, 3, 3, 3)).astype(np.float32), name="W")
|
|
conv = helper.make_node("Conv", ["X", "W"], ["C"], kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
pool = helper.make_node("MaxPool", ["C"], ["Y"], kernel_shape=[2, 2], strides=[2, 2], pads=[0, 0, 0, 0])
|
|
graph = helper.make_graph([conv, pool], "maxpool_after_conv", [X], [Y], initializer=[W])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def relu_basic():
|
|
"""Standalone Relu 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("Relu", ["X"], ["Y"])
|
|
graph = helper.make_graph([node], "relu_basic", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "relu/basic", "relu_basic.onnx")
|
|
|
|
|
|
def relu_4d():
|
|
"""Standalone Relu 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("Relu", ["X"], ["Y"])
|
|
graph = helper.make_graph([node], "relu_4d", [X], [Y])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "relu/4d", "relu_4d.onnx")
|
|
|
|
|
|
def relu_after_conv():
|
|
"""Conv followed by Relu."""
|
|
rng = np.random.default_rng(60)
|
|
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")
|
|
conv = helper.make_node("Conv", ["X", "W", "B"], ["C"],
|
|
kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
|
|
relu = helper.make_node("Relu", ["C"], ["Y"])
|
|
graph = helper.make_graph([conv, relu], "relu_after_conv", [X], [Y], initializer=[W, B])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
save_model(model, "relu/after_conv", "relu_after_conv.onnx")
|
|
|
|
|
|
def relu_after_gemm():
|
|
"""Gemm followed by Relu."""
|
|
B, K, N = 4, 64, 32
|
|
rng = np.random.default_rng(61)
|
|
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"])
|
|
relu = helper.make_node("Relu", ["G"], ["Y"])
|
|
graph = helper.make_graph([gemm, relu], "relu_after_gemm", [A], [Y], initializer=[W, C])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
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
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
print("Generating GEMM tests:")
|
|
gemm_non_square()
|
|
gemm_with_bias()
|
|
gemm_transB()
|
|
gemm_alpha_beta()
|
|
gemm_small()
|
|
gemm_large()
|
|
gemm_transB_with_bias()
|
|
|
|
print("\nGenerating Conv tests:")
|
|
conv_3x3_kernel()
|
|
conv_stride2()
|
|
conv_multi_channel()
|
|
conv_1x1()
|
|
conv_same_padding_3x3()
|
|
conv_explicit_padding()
|
|
conv_with_bias_3x3()
|
|
conv_batch_2()
|
|
conv_large_spatial()
|
|
|
|
print("\nGenerating Pooling tests:")
|
|
maxpool_basic()
|
|
maxpool_stride2_multichannel()
|
|
maxpool_same_upper()
|
|
avgpool_basic()
|
|
avgpool_explicit_padding()
|
|
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.")
|