diff --git a/onnxconverter_common/onnx_ops.py b/onnxconverter_common/onnx_ops.py index 1eea627..a5cda27 100644 --- a/onnxconverter_common/onnx_ops.py +++ b/onnxconverter_common/onnx_ops.py @@ -336,9 +336,25 @@ def apply_floor(scope, input_name, output_name, container, operator_name=None): _apply_unary_operation(scope, 'Floor', input_name, output_name, container, operator_name=operator_name) -def apply_flatten(scope, input_name, output_name, container, operator_name=None): +def apply_flatten(scope, input_name, output_name, container, operator_name=None, axis=1): name = _create_name_or_use_existing_one(scope, 'Flatten', operator_name) - container.add_node('Flatten', input_name, output_name, name=name) + if container.target_opset < 9: + op_version = 1 + elif container.target_opset < 11: + op_version = 9 + else: + op_version = 11 + container.add_node('Flatten', input_name, output_name, name=name, op_version=op_version, axis=axis) + + +def apply_gather(scope, input_names, output_name, container, operator_name=None, axis=0): + name = _create_name_or_use_existing_one(scope, 'Gather', operator_name) + if container.target_opset < 11: + op_version = 1 + else: + op_version = 11 + + container.add_node('Gather', input_names, output_name, name=name, op_version=op_version, axis=axis) def apply_gemm(scope, input_name, output_name, container, operator_name=None, alpha=1.0, beta=1.0, @@ -362,6 +378,34 @@ def apply_gemm(scope, input_name, output_name, container, operator_name=None, al container.add_node('Gemm', input_name, output_name, name=name, **attrs) +def apply_greater(scope, input_names, output_name, container, operator_name=None): + name = _create_name_or_use_existing_one(scope, 'Greater', operator_name) + if container.target_opset < 7: + op_version = 1 + elif container.target_opset < 9: + op_version = 7 + else: + op_version = 9 + + container.add_node('Greater', input_names, output_name, name=name, op_version=op_version) + + +def apply_gru(scope, input_names, output_names, container, operator_name=None, output_seq=0, reset_after=0, **attrs): + name = _create_name_or_use_existing_one(scope, 'GRU', operator_name) + if container.target_opset < 3: + op_version = 1 + attrs['output_sequence'] = 1 if output_seq else 0 + else: + attrs['linear_before_reset'] = 1 if reset_after else 0 + if container.target_opset <= 5: + attrs['output_sequence'] = 1 if output_seq else 0 + op_version = 3 + else: + op_version = 7 + + container.add_node('GRU', input_names, output_names, name=name, op_version=op_version, **attrs) + + def apply_hard_sigmoid(scope, input_name, output_name, container, operator_name=None, alpha=None, beta=None): _apply_unary_operation(scope, 'HardSigmoid', input_name, output_name, container, operator_name, alpha=alpha, beta=beta) @@ -393,10 +437,24 @@ def apply_log(scope, input_name, output_name, container, operator_name=None): _apply_unary_operation(scope, 'Log', input_name, output_name, container, operator_name=operator_name) +def apply_lstm(scope, input_names, output_names, container, operator_name=None, output_seq=0, **attrs): + name = _create_name_or_use_existing_one(scope, 'LSTM', operator_name) + if container.target_opset <= 6: + attrs['output_sequence'] = 1 if output_seq else 0 + op_version = 1 + else: + op_version = 7 + container.add_node('LSTM', input_names, output_names, name=name, op_version=op_version, **attrs) + + def apply_matmul(scope, input_names, output_name, container, operator_name=None): op_type = 'MatMul' name = _create_name_or_use_existing_one(scope, op_type, operator_name) - container.add_node(op_type, input_names, output_name, op_version=9, name=name) + if container.target_opset <= 9: + op_version = 1 + else: + op_version = 9 + container.add_node(op_type, input_names, output_name, op_version=op_version, name=name) def apply_max(scope, input_names, output_name, container, operator_name=None): @@ -425,7 +483,8 @@ def apply_normalization(scope, input_name, output_name, container, operator_name container.add_node('LpNormalization', input_name, output_name, name=name, p=p, axis=axis) -def apply_pad(scope, input_name, output_name, container, operator_name=None, mode=None, pads=None, value=None): +def apply_pad(scope, input_name, output_name, container, operator_name=None, mode=None, pads=None, value=None, + onnx_type=onnx_proto.TensorProto.FLOAT): name = _create_name_or_use_existing_one(scope, 'Pad', operator_name) attrs = {'name': name} inputs = input_name if isinstance(input_name, list) else [input_name] @@ -449,7 +508,7 @@ def apply_pad(scope, input_name, output_name, container, operator_name=None, mod inputs.append(pads_name) if value is not None: value_name = scope.get_unique_variable_name(name + '_value') - container.add_initializer(value_name, onnx_proto.TensorProto.FLOAT, [], [value]) + container.add_initializer(value_name, onnx_type, [], [value]) inputs.append(value_name) container.add_node('Pad', inputs, output_name, op_version=op_version, **attrs) @@ -602,6 +661,16 @@ def apply_resize(scope, input_name, output_name, container, operator_name=None, container.add_node('Resize', inputs, output_name, op_version=op_version, **attrs) +def apply_rnn(scope, input_names, output_names, container, operator_name=None, output_seq=0, **attrs): + name = _create_name_or_use_existing_one(scope, 'RNN', operator_name) + if container.target_opset <= 6: + attrs['output_sequence'] = 1 if output_seq else 0 + op_version = 1 + else: + op_version = 7 + container.add_node('RNN', input_names, output_names, name=name, op_version=op_version, **attrs) + + def apply_sigmoid(scope, input_name, output_name, container, operator_name=None): _apply_unary_operation(scope, 'Sigmoid', input_name, output_name, container, operator_name) @@ -843,3 +912,14 @@ def apply_upsample(scope, input_name, output_name, container, operator_name=None # We implement Upsample through Resize instead apply_resize(scope, input_name, output_name, container, operator_name, mode, coordinate_transformation_mode, scales) + + +def apply_unsqueeze(scope, input_name, output_name, container, operator_name=None, axis=0, rank=0): + name = _create_name_or_use_existing_one(scope, 'Unsqueeze', operator_name) + if container.target_opset < 11: + op_version = 1 + if axis < 0: + axis += rank + 1 + else: + op_version = 11 + container.add_node('Unsqueeze', input_name, output_name, name=name, op_version=op_version, axes=[axis])