def test_make_optional_value_info(self): # type: () -> None
tensor_type_proto = helper.make_tensor_type_proto(elem_type=2,
shape=[5])
tensor_val_into = helper.make_value_info(name='test',
type_proto=tensor_type_proto)
optional_type_proto = helper.make_optional_type_proto(
tensor_type_proto)
optional_val_info = helper.make_value_info(
name='test', type_proto=optional_type_proto)
self.assertEqual(optional_val_info.name, 'test')
self.assertTrue(optional_val_info.type.optional_type)
self.assertEqual(optional_val_info.type.optional_type.elem_type,
tensor_val_into.type)
# Test Sequence
sequence_type_proto = helper.make_sequence_type_proto(
tensor_type_proto)
optional_type_proto = helper.make_optional_type_proto(
sequence_type_proto)
optional_val_info = helper.make_value_info(
name='test', type_proto=optional_type_proto)
self.assertEqual(optional_val_info.name, 'test')
self.assertTrue(optional_val_info.type.optional_type)
sequence_value_info = helper.make_value_info(
name='test', type_proto=tensor_type_proto)
self.assertEqual(
optional_val_info.type.optional_type.elem_type.sequence_type.
elem_type, sequence_value_info.type)
def test_make_seuence_value_info(self): # type: () -> None
tensor_type_proto = helper.make_tensor_type_proto(elem_type=2,
shape=None)
sequence_type_proto = helper.make_sequence_type_proto(
tensor_type_proto)
sequence_val_info = helper.make_value_info(
name='test', type_proto=sequence_type_proto)
sequence_val_info_prim = helper.make_tensor_sequence_value_info(
name='test', elem_type=2, shape=None)
self.assertEqual(sequence_val_info, sequence_val_info_prim)
def _test_op_upgrade(
self,
op: Text,
from_opset: int,
input_shapes: List[Union[List[Optional[int]], Text]] = [[3, 4, 5]],
output_shapes: List[List[Optional[int]]] = [[3, 4, 5]],
input_types: Union[List[Any], None] = None,
output_types: Union[List[Any], None] = None,
initializer: List[Any] = [],
attrs: Dict[Text, Any] = {},
seq_inputs: List[int] = [],
seq_outputs: List[int] = [],
optional_inputs: List[int] = [],
optional_outputs: List[int] = []
) -> None:
global tested_ops
tested_ops.append(op)
n_inputs = len(input_shapes)
letters = list(string.ascii_lowercase)[:n_inputs]
input_names = [
letter if shape != '' else '' for (letter, shape) in zip(letters, input_shapes)
]
if input_types is None:
input_types = [TensorProto.FLOAT] * n_inputs
is_sequence = [0 if id not in seq_inputs else 1 for id in range(n_inputs)]
is_optional = [0 if id not in optional_inputs else 1 for id in range(n_inputs)]
# turn empty strings into [0] to ease type analysis, even though those entries
# will be ignored
input_shapes_cast = cast(List[List[int]],
[[0] if isinstance(shape, str) else shape for shape in input_shapes]
)
inputs: List[ValueInfoProto] = []
for (name, ttype, shape, is_seq, is_opt) in \
zip(input_names, input_types, input_shapes_cast, is_sequence, is_optional):
if name != '':
if is_seq:
inputs += [helper.make_tensor_sequence_value_info(name, ttype, shape)]
elif is_opt:
type_proto = helper.make_tensor_type_proto(ttype, shape)
optional_type_proto = helper.make_optional_type_proto(type_proto)
inputs += [helper.make_value_info(name, optional_type_proto)]
else:
inputs += [helper.make_tensor_value_info(name, ttype, shape)]
n_outputs = len(output_shapes)
output_names = list(string.ascii_lowercase)[n_inputs:n_inputs + n_outputs]
if output_types is None:
output_types = [TensorProto.FLOAT] * n_outputs
is_sequence = [0 if id not in seq_outputs else 1 for id in range(n_outputs)]
is_optional = [0 if id not in optional_outputs else 1 for id in range(n_outputs)]
output_shapes_cast = cast(List[List[int]],
[[0] if isinstance(shape, str) else shape for shape in output_shapes]
)
outputs: List[ValueInfoProto] = []
for (name, ttype, shape, is_seq, is_opt) in \
zip(output_names, output_types, output_shapes_cast, is_sequence, is_optional):
if is_seq:
outputs += [helper.make_tensor_sequence_value_info(name, ttype, shape)]
elif is_opt:
type_proto = helper.make_tensor_type_proto(ttype, shape)
optional_type_proto = helper.make_optional_type_proto(type_proto)
outputs += [helper.make_value_info(name, optional_type_proto)]
else:
outputs += [helper.make_tensor_value_info(name, ttype, shape)]
node = helper.make_node(op, input_names, output_names, **attrs)
graph = helper.make_graph([node], op, inputs, outputs, initializer)
original = helper.make_model(
graph,
producer_name='test',
opset_imports=[helper.make_opsetid('', from_opset)]
)
onnx.checker.check_model(original)
shape_inference.infer_shapes(original, strict_mode=True)
converted = version_converter.convert_version(original, latest_opset)
onnx.checker.check_model(converted)
shape_inference.infer_shapes(converted, strict_mode=True)
def test_aten_embedding(self):
class NeuralNetEmbedding(torch.nn.Module):
def __init__(self, num_embeddings, embedding_dim, hidden_size):
super(NeuralNetEmbedding, self).__init__()
self.embedding = torch.nn.Embedding(num_embeddings,
embedding_dim)
self.linear = torch.nn.Linear(embedding_dim, hidden_size)
def forward(self, input):
embedding_result = self.embedding(input)
return embedding_result, self.linear(embedding_result)
N, num_embeddings, embedding_dim, hidden_size = 64, 32, 128, 128
model = NeuralNetEmbedding(num_embeddings, embedding_dim, hidden_size)
with torch.no_grad():
x = torch.randint(high=num_embeddings,
size=(N, ),
dtype=torch.int64)
dynamic_axes = {
"x": {
0: "x_dim0"
},
"y": {
0: "y_dim0",
1: "y_dim1"
}
}
f = io.BytesIO()
export(
model,
x,
f=f,
input_names=["x"],
output_names=["y"],
dynamic_axes=dynamic_axes,
opset_version=14,
)
exported_model = onnx.load_model_from_string(f.getvalue())
# PyTorch exporter emitting ATen op is still under development. Currently convert it manually for testing.
for node in exported_model.graph.node:
if node.op_type == "Gather":
node.domain = "org.pytorch.aten"
node.op_type = "ATen"
attr = node.attribute.add()
attr.name = "operator"
attr.type = 3
attr.s = "embedding".encode()
exported_model.graph.node.append(
helper.make_node(
"Constant",
[],
["padding_idx"],
value=helper.make_tensor("padding_idx",
TensorProto.INT64, (),
[-1]),
))
exported_model.graph.node.append(
helper.make_node(
"Constant",
[],
["scale_grad_by_freq"],
value=helper.make_tensor("scale_grad_by_freq",
TensorProto.BOOL, (),
[False]),
))
exported_model.graph.node.append(
helper.make_node(
"Constant",
[],
["sparse"],
value=helper.make_tensor("sparse",
TensorProto.BOOL, (),
[False]),
))
node.input.append("padding_idx")
node.input.append("scale_grad_by_freq")
node.input.append("sparse")
exported_model.graph.value_info.append(
helper.make_value_info(
name=node.output[0],
type_proto=helper.make_tensor_type_proto(
elem_type=TensorProto.FLOAT,
shape=[
node.output[0] + "_dim0",
node.output[0] + "_dim1"
]),
))
break
# The ONNX graph to run contains ATen Op.
assert any(node.op_type == "ATen"
for node in exported_model.graph.node)
init_aten_op_executor()
# Run w/o IO binding.
for _ in range(8):
x = torch.randint(high=num_embeddings,
size=(N, ),
dtype=torch.int64)
pt_y1, pt_y2 = model(x)
session = ort.InferenceSession(exported_model.SerializeToString(),
providers=["CPUExecutionProvider"])
ort_y1, ort_y2 = session.run([], {"x": x.numpy()})
np.testing.assert_almost_equal(ort_y1, pt_y1.detach().numpy())
np.testing.assert_almost_equal(ort_y2, pt_y2.detach().numpy())
# Run w/ IO binding.
for _ in range(8):
x = torch.randint(high=num_embeddings,
size=(N, ),
dtype=torch.int64)
ort_x = ort.OrtValue.ortvalue_from_numpy(x.detach().numpy(), "cpu")
pt_y1, pt_y2 = model(x)
np_y1 = np.zeros(tuple(pt_y1.size()), dtype=np.float32)
np_y2 = np.zeros(tuple(pt_y2.size()), dtype=np.float32)
ort_y1 = ort.OrtValue.ortvalue_from_numpy(np_y1, "cpu")
ort_y2 = ort.OrtValue.ortvalue_from_numpy(np_y2, "cpu")
session = ort.InferenceSession(exported_model.SerializeToString(),
providers=["CPUExecutionProvider"])
io_binding = session.io_binding()
io_binding.bind_ortvalue_input(exported_model.graph.input[0].name,
ort_x)
io_binding.bind_ortvalue_output(
exported_model.graph.output[0].name, ort_y1)
io_binding.bind_ortvalue_output(
exported_model.graph.output[1].name, ort_y2)
session.run_with_iobinding(io_binding)
np.testing.assert_almost_equal(np_y1, pt_y1.detach().numpy())
np.testing.assert_almost_equal(np_y2, pt_y2.detach().numpy())
