Skip to content

model and test data to fix einsum importer #1132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Dec 25, 2023
Merged

model and test data to fix einsum importer #1132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
81ad54e
undo comment
fengyuentau Dec 22, 2023
a8827e0
add model and data
fengyuentau Dec 22, 2023
85d8061
update script
fengyuentau Dec 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Binary file added testdata/dnn/onnx/data/input_einsum_const_inputs.npy
View file
Open in desktop
Binary file not shown.
Binary file added testdata/dnn/onnx/data/output_einsum_const_inputs.npy
View file
Open in desktop
Binary file not shown.
206 changes: 109 additions & 97 deletions testdata/dnn/onnx/generate_onnx_models_with_onnxscript.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,103 +43,103 @@ def make_model_and_data(model, *args, **kwargs):
output_files = os.path.join("data", "output_" + name)
np.save(output_files, output)

# '''
# It builds a model with two Gather ops sharing a single same indices:

# [Input] -> Gather(indices=0) -> Gather(indices=0) -> [Output]

# , where the two indices constants have the same name.
# '''
# @ost.script()
# def gather_shared_indices(x: ost.FLOAT[2, 1, 3, 4]) -> ost.FLOAT[3, 4]:
# indices = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([0], dtype=np.int64)))
# y0 = op.Gather(x, indices, axis=0)
# y1 = op.Gather(y0, indices, axis=0)
# return y1
# make_model_and_data(gather_shared_indices, np.random.rand(2, 1, 3, 4).astype(np.float32))

# '''
# [Input] -> Greater(B=61) -> [Output]
# \
# dtype=np.int64
# '''
# @ost.script()
# def greater_input_dtype_int64(x: ost.FLOAT[27, 9]) ->ost.BOOL[27, 9]:
# y = op.Greater(x, op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([61], dtype=np.int64))))
# return y
# make_model_and_data(greater_input_dtype_int64, np.random.randint(0, 100, size=[27, 9], dtype=np.int64), force_saving_input_as_dtype_float32=True, force_saving_output_as_dtype_float32=True)

# from onnxscript import opset11

# @ost.script()
# def two_resizes_with_shared_subgraphs(x: ost.FLOAT["batch", 1, "height", "width"], y: ost.FLOAT[1, 1, 3, 2], z: ost.FLOAT[1, 1, 2, 1]) ->ost.FLOAT["batch", 1, "height", "width"]:
# shape_src_1 = opset11.Shape(x)
# shape_src_2 = opset11.Shape(x)
# gather_h = opset11.Gather(shape_src_1, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([2], dtype=np.int64))), axis=0)
# gather_w = opset11.Gather(shape_src_2, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([3], dtype=np.int64))), axis=0)
# unsqueeze_w_1 = opset11.Unsqueeze(gather_w, axes=[0])
# unsqueeze_w_2 = opset11.Unsqueeze(gather_w, axes=[0])
# unsqueeze_h_1 = opset11.Unsqueeze(gather_h, axes=[0])
# unsqueeze_h_2 = opset11.Unsqueeze(gather_h, axes=[0])
# concat_1 = opset11.Cast(opset11.Concat(unsqueeze_h_1, unsqueeze_w_1, axis=0), to=ost.INT64.dtype)
# concat_2 = opset11.Cast(opset11.Concat(unsqueeze_h_2, unsqueeze_w_2, axis=0), to=ost.INT64.dtype)

# # This op is required to test double node removal
# y = opset11.Add(y, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [1], np.array([0.5], dtype=np.float32))))

# # First branch
# sliced = opset11.Slice(opset11.Shape(y),
# starts=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
# ends=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([2], dtype=np.int64))),
# axes=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
# )
# concat_y = opset11.Concat(sliced, concat_1, axis=0)
# resized_y = opset11.Resize(y,
# roi=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
# scales=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
# sizes=concat_y,
# coordinate_transformation_mode='pytorch_half_pixel',
# cubic_coeff_a=-0.75,
# mode='linear',
# nearest_mode='floor'
# )

# # Second branch
# sliced = opset11.Slice(opset11.Shape(z),
# starts=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
# ends=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([2], dtype=np.int64))),
# axes=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
# )
# concat_z = opset11.Concat(sliced, concat_2, axis=0)
# resized_z = opset11.Resize(z,
# roi=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
# scales=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
# sizes=concat_z,
# coordinate_transformation_mode='pytorch_half_pixel',
# cubic_coeff_a=-0.75,
# mode='linear',
# nearest_mode='floor'
# )

# return opset11.Add(resized_y, resized_z)

# make_model_and_data(two_resizes_with_shared_subgraphs, np.random.rand(1, 1, 4, 5).astype(np.float32), np.random.rand(1, 1, 3, 2).astype(np.float32), np.random.rand(1, 1, 2, 1).astype(np.float32))


# @ost.script()
# def bias_gelu(x: ost.FLOAT[1, 2, 3]) -> ost.FLOAT[1, 2, 3]:
# bias = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [3], np.array([0.1, 0.3, 0.2], dtype=np.float32)))
# add1 = op.Add(x, bias)
# tmp = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([np.sqrt(2)], dtype=np.float32)))
# div = op.Div(add1, tmp)
# erf = op.Erf(div)
# tmp_0 = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([1], dtype=np.float32)))
# add2 = op.Add(erf, tmp_0)
# mul = op.Mul(add1, add2)
# tmp_1 = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([0.5], dtype=np.float32)))
# return op.Mul(mul, tmp_1)

# make_model_and_data(bias_gelu, np.random.rand(1, 2, 3).astype(np.float32))
'''
It builds a model with two Gather ops sharing a single same indices:

[Input] -> Gather(indices=0) -> Gather(indices=0) -> [Output]

, where the two indices constants have the same name.
'''
@ost.script()
def gather_shared_indices(x: ost.FLOAT[2, 1, 3, 4]) -> ost.FLOAT[3, 4]:
indices = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([0], dtype=np.int64)))
y0 = op.Gather(x, indices, axis=0)
y1 = op.Gather(y0, indices, axis=0)
return y1
make_model_and_data(gather_shared_indices, np.random.rand(2, 1, 3, 4).astype(np.float32))

'''
[Input] -> Greater(B=61) -> [Output]
\
dtype=np.int64
'''
@ost.script()
def greater_input_dtype_int64(x: ost.FLOAT[27, 9]) ->ost.BOOL[27, 9]:
y = op.Greater(x, op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([61], dtype=np.int64))))
return y
make_model_and_data(greater_input_dtype_int64, np.random.randint(0, 100, size=[27, 9], dtype=np.int64), force_saving_input_as_dtype_float32=True, force_saving_output_as_dtype_float32=True)

from onnxscript import opset11

@ost.script()
def two_resizes_with_shared_subgraphs(x: ost.FLOAT["batch", 1, "height", "width"], y: ost.FLOAT[1, 1, 3, 2], z: ost.FLOAT[1, 1, 2, 1]) ->ost.FLOAT["batch", 1, "height", "width"]:
shape_src_1 = opset11.Shape(x)
shape_src_2 = opset11.Shape(x)
gather_h = opset11.Gather(shape_src_1, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([2], dtype=np.int64))), axis=0)
gather_w = opset11.Gather(shape_src_2, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [], np.array([3], dtype=np.int64))), axis=0)
unsqueeze_w_1 = opset11.Unsqueeze(gather_w, axes=[0])
unsqueeze_w_2 = opset11.Unsqueeze(gather_w, axes=[0])
unsqueeze_h_1 = opset11.Unsqueeze(gather_h, axes=[0])
unsqueeze_h_2 = opset11.Unsqueeze(gather_h, axes=[0])
concat_1 = opset11.Cast(opset11.Concat(unsqueeze_h_1, unsqueeze_w_1, axis=0), to=ost.INT64.dtype)
concat_2 = opset11.Cast(opset11.Concat(unsqueeze_h_2, unsqueeze_w_2, axis=0), to=ost.INT64.dtype)

# This op is required to test double node removal
y = opset11.Add(y, opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [1], np.array([0.5], dtype=np.float32))))

# First branch
sliced = opset11.Slice(opset11.Shape(y),
starts=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
ends=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([2], dtype=np.int64))),
axes=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
)
concat_y = opset11.Concat(sliced, concat_1, axis=0)
resized_y = opset11.Resize(y,
roi=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
scales=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
sizes=concat_y,
coordinate_transformation_mode='pytorch_half_pixel',
cubic_coeff_a=-0.75,
mode='linear',
nearest_mode='floor'
)

# Second branch
sliced = opset11.Slice(opset11.Shape(z),
starts=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
ends=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([2], dtype=np.int64))),
axes=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.INT64, [1], np.array([0], dtype=np.int64))),
)
concat_z = opset11.Concat(sliced, concat_2, axis=0)
resized_z = opset11.Resize(z,
roi=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
scales=opset11.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [0], np.empty([0]))),
sizes=concat_z,
coordinate_transformation_mode='pytorch_half_pixel',
cubic_coeff_a=-0.75,
mode='linear',
nearest_mode='floor'
)

return opset11.Add(resized_y, resized_z)

make_model_and_data(two_resizes_with_shared_subgraphs, np.random.rand(1, 1, 4, 5).astype(np.float32), np.random.rand(1, 1, 3, 2).astype(np.float32), np.random.rand(1, 1, 2, 1).astype(np.float32))


@ost.script()
def bias_gelu(x: ost.FLOAT[1, 2, 3]) -> ost.FLOAT[1, 2, 3]:
bias = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [3], np.array([0.1, 0.3, 0.2], dtype=np.float32)))
add1 = op.Add(x, bias)
tmp = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([np.sqrt(2)], dtype=np.float32)))
div = op.Div(add1, tmp)
erf = op.Erf(div)
tmp_0 = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([1], dtype=np.float32)))
add2 = op.Add(erf, tmp_0)
mul = op.Mul(add1, add2)
tmp_1 = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, [], np.array([0.5], dtype=np.float32)))
return op.Mul(mul, tmp_1)

make_model_and_data(bias_gelu, np.random.rand(1, 2, 3).astype(np.float32))
Comment on lines +46 to +142
Copy link
Contributor

@asmorkalov asmorkalov Dec 25, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@fengyuentau It's not relevant change.


batch_size = 1
sequence_length = 320
Expand Down Expand Up @@ -298,3 +298,15 @@ def attention_single_head(x: ost.FLOAT[batch_size, sequence_length, input_hidden
return reshape

make_model_and_data(attention_single_head, np.random.rand(batch_size, sequence_length, input_hidden_size).astype(np.float32))

# Einsum_const_inputs

input_0_data = np.random.rand(3, 2, 2, 4).astype(np.float32)
input_1_data = np.random.rand(2, 2, 4).astype(np.float32)

@ost.script()
def einsum_const_inputs(input_0: ost.FLOAT[3, 2, 2, 4]) -> ost.FLOAT[3, 2, 2, 2]:
input_1 = op.Constant(value=onnx.helper.make_tensor("", onnx.TensorProto.FLOAT, input_1_data.shape, input_1_data))
return op.Einsum(input_0, input_1, equation="bhwc, hkc -> bhwk")

make_model_and_data(einsum_const_inputs, input_0_data)
Binary file added testdata/dnn/onnx/models/einsum_const_inputs.onnx
View file
Open in desktop
Binary file not shown.