def generic_convert_binary1_or_binary2(converter, nnef_op, onnx_graph,
target_name1, target_name2):
# type: (Converter, NNEFOperation, ONNXGraph, str, str)->None
x, y = converter.converted_tensors(nnef_op.inputs)
z = converter.converted_tensor(nnef_op.output)
binary1 = ONNXOperation(
graph=onnx_graph,
name=target_name1,
inputs=(
converter.add_unsqueeze(onnx_graph, x, list(range(x.rank, y.rank)))
if 0 < x.rank < y.rank else x,
converter.add_unsqueeze(onnx_graph, y, list(range(y.rank, x.rank)))
if 0 < y.rank < x.rank else y),
outputs=ONNXTensor(graph=onnx_graph,
shape=list(z.shape),
dtype=z.dtype))
binary2 = ONNXOperation(
graph=onnx_graph,
name=target_name2,
inputs=(
converter.add_unsqueeze(onnx_graph, x, list(range(x.rank, y.rank)))
if 0 < x.rank < y.rank else x,
converter.add_unsqueeze(onnx_graph, y, list(range(y.rank, x.rank)))
if 0 < y.rank < x.rank else y),
outputs=ONNXTensor(graph=onnx_graph,
shape=list(z.shape),
dtype=z.dtype))
ONNXOperation(graph=onnx_graph,
name="Or",
inputs=(binary1.output, binary2.output),
outputs=z)
def convert_unstack(converter, nnef_op, onnx_graph):
# type: (Converter, NNEFOperation, ONNXGraph)->None
input = converter.converted_tensor(nnef_op.input)
outputs = converter.converted_tensors(nnef_op.outputs)
axis = nnef_op.attribs['axis']
num_parts = input.shape[axis]
def unsqueeze(shape, axis):
shape = list(shape)
shape.insert(axis, 1)
return shape
intermediates = [
ONNXTensor(graph=onnx_graph,
shape=unsqueeze(output.shape, axis),
dtype=output.dtype) for output in outputs
]
ONNXOperation(graph=onnx_graph,
name='Split',
inputs=input,
attribs=dict(axis=axis, split=[1] * num_parts),
outputs=intermediates)
for intermediate, output in zip(intermediates, outputs):
converter.add_unsqueeze(onnx_graph,
intermediate,
axis,
unsqueezed_tensor=output)
def generic_convert_binary(converter,
nnef_op,
onnx_graph,
target_name,
broadcast_workaround=False):
# type: (Converter, NNEFOperation, ONNXGraph, str, bool)->None
if broadcast_workaround:
x, y = converter.converted_tensors(nnef_op.inputs)
z = converter.converted_tensor(nnef_op.output)
# Caffe2 onnx backend cannot broadcast in Pow
if x.shape != z.shape:
x = converter.add_expand(onnx_graph, x, z.shape)
if y.shape != z.shape:
y = converter.add_expand(onnx_graph, y, z.shape)
ONNXOperation(graph=onnx_graph,
name=target_name,
inputs=(x, y),
outputs=z)
else:
x, y = converter.converted_tensors(nnef_op.inputs)
z = converter.converted_tensor(nnef_op.output)
ONNXOperation(graph=onnx_graph,
name=target_name,
inputs=(converter.add_unsqueeze(
onnx_graph, x, list(range(x.rank, y.rank)))
if 0 < x.rank < y.rank else x,
converter.add_unsqueeze(
onnx_graph, y, list(range(y.rank, x.rank)))
if 0 < y.rank < x.rank else y),
outputs=z)
def convert_batch_normalization(converter, nnef_op, onnx_graph):
# type: (Converter, NNEFOperation, ONNXGraph)->None
input, mean, variance, offset, scale = converter.converted_tensors(
nnef_op.inputs)
output = converter.converted_tensor(nnef_op.output)
if input.rank < 3: # Caffe2 BatchNorm works only for rank >= 3
axes = list(range(input.rank, 3))
input = converter.add_unsqueeze(onnx_graph=onnx_graph,
onnx_tensor=input,
axes=axes)
final_output = output
output = ONNXTensor(graph=onnx_graph,
shape=list(input.shape),
dtype=input.dtype)
converter.add_squeeze(onnx_graph=onnx_graph,
onnx_tensor=output,
axes=axes,
squeezed_tensor=final_output)
ONNXOperation(graph=onnx_graph,
name='BatchNormalization',
inputs=(input, converter.add_squeeze(onnx_graph, scale, [0]),
converter.add_squeeze(onnx_graph, offset, [0]),
converter.add_squeeze(onnx_graph, mean, [0]),
converter.add_squeeze(onnx_graph, variance, [0])),
attribs=dict(epsilon=nnef_op.attribs['epsilon']),
outputs=output)
def convert_ne(converter, nnef_op, onnx_graph):
# type: (Converter, NNEFOperation, ONNXGraph)->None
x, y = converter.converted_tensors(nnef_op.inputs)
z = converter.converted_tensor(nnef_op.output)
equal = ONNXOperation(
graph=onnx_graph,
name='Equal',
inputs=(converter.add_unsqueeze(onnx_graph, x, list(range(x.rank, y.rank))) if 0 < x.rank < y.rank else x,
converter.add_unsqueeze(onnx_graph, y, list(range(y.rank, x.rank))) if 0 < y.rank < x.rank else y),
outputs=ONNXTensor(graph=onnx_graph, shape=list(z.shape), dtype=z.dtype))
ONNXOperation(
graph=onnx_graph,
name="Not",
inputs=equal.output,
outputs=z)
def convert_stack(converter, nnef_op, onnx_graph):
# type: (Converter, NNEFOperation, ONNXGraph)->None
inputs = converter.converted_tensors(nnef_op.inputs)
output = converter.converted_tensor(nnef_op.output)
inputs = [converter.add_unsqueeze(onnx_graph, input, [nnef_op.attribs['axis']]) for input in inputs]
ONNXOperation(graph=onnx_graph,
name='Concat',
inputs=inputs,
attribs=dict(axis=nnef_op.attribs['axis']),
outputs=output)
def broadcast(t):
return converter.add_unsqueeze(
onnx_graph, t, list(range(
t.rank, max_rank))) if 0 < t.rank < max_rank else t
