Merge pull request #1114 from xadupre/fft

xadupre · web-flow · commit b51df2fa8943 · 2020-09-29T10:27:38.000+02:00
Add support for operators RFFT, ComplexAbs
diff --git a/tests/test_backend.py b/tests/test_backend.py
@@ -24,6 +24,7 @@
 from tf2onnx import constants, utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tf_loader import is_tf2
+from tf2onnx.onnx_opset.signal import make_dft_constant
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,function-redefined,cell-var-from-loop
 
@@ -3624,6 +3625,27 @@ def test_conv2d_1_kernel_as_input(self):
                           [1., 1., 4.]], dtype=np.float32).reshape(_KERNEL3x3)
         self._conv_kernel_as_input_test(x_val, w_val)
 
+    @check_tf_min_version("1.14")
+    def test_rfft_ops(self):
+
+        def dft_slow(x, M):
+            xt = x.T
+            res = np.dot(M, xt)
+            return np.transpose(res, (0, 2, 1))
+
+        x_val = make_xval([2, 4]).astype(np.float32)
+        M_both = make_dft_constant(x_val.shape[1], x_val.dtype, x_val.shape[1])
+        fft = dft_slow(x_val, M_both)
+        fft_npy = np.fft.rfft(x_val)
+        assert_almost_equal(fft[0, :, :], np.real(fft_npy))
+        assert_almost_equal(fft[1, :, :], np.imag(fft_npy))
+
+        x_val = make_xval([3, 4]).astype(np.float32)
+        def func(x):
+            op_ = tf.signal.rfft(x)
+            return tf.abs(op_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
 
 if __name__ == '__main__':
     unittest_main()
diff --git a/tf2onnx/custom_opsets/ms.py b/tf2onnx/custom_opsets/ms.py
@@ -23,7 +23,9 @@ def make_range(ctx, start, limit, delta, output, scope_name, shape, dtype):
 
 def _make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dtype):
     utils.make_sure(
-        dtype in [TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.INT16, TensorProto.INT32, TensorProto.INT64],
+        dtype in [TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.INT16,
+                  TensorProto.INT32, TensorProto.INT64,
+                  TensorProto.COMPLEX64, TensorProto.COMPLEX128],
         "dtype %s is not supported", dtype)
     ctx.make_node("Range", [start, limit, delta], outputs=[output], name=scope_name, shapes=[shape], dtypes=[dtype],
                   domain=constants.MICROSOFT_DOMAIN)
diff --git a/tf2onnx/graph.py b/tf2onnx/graph.py
@@ -988,7 +988,7 @@ def _get_unvisited_child(g, node, not_visited):
             all_input = list(filter(lambda a: a != '', all_input))
             for inp in sorted(all_input):
                 j = self.get_node_by_output(inp)
-                utils.make_sure(j is not None, "Cannot find node with output {}".format(inp))
+                utils.make_sure(j is not None, "Cannot find node with output %r", inp)
                 if self.parent_graph and j.name not in op_name_to_index:
                     # there might be some outer-scoped inputs for an inner Graph.
                     pass
diff --git a/tf2onnx/onnx_opset/__init__.py b/tf2onnx/onnx_opset/__init__.py
@@ -2,4 +2,18 @@
 # Licensed under the MIT license.
 """tf2onnx.onnx_opset module"""
 
-from . import common, controlflow, generator, logical, math, misc, nn, quantize, reduction, rnn, tensor, traditionalml
+from . import (
+    common,
+    controlflow,
+    generator,
+    logical,
+    math,
+    misc,
+    nn,
+    quantize,
+    reduction,
+    rnn,
+    signal,
+    tensor,
+    traditionalml
+)
diff --git a/tf2onnx/onnx_opset/signal.py b/tf2onnx/onnx_opset/signal.py
@@ -0,0 +1,228 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+signal
+"""
+
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import logging
+
+import numpy as np
+from onnx import onnx_pb
+from onnx.numpy_helper import to_array
+from tf2onnx import utils
+from tf2onnx.handler import tf_op
+
+logger = logging.getLogger(__name__)
+
+
+# pylint: disable=unused-argument,missing-docstring
+
+def make_dft_constant(length, dtype, fft_length):
+    n = np.arange(length)
+    k = n.reshape((length, 1)).astype(np.float64)
+    mat = np.exp(-2j * np.pi * k * n / length)
+    mat = mat[:fft_length // 2 + 1]
+    both = np.empty((2,) + mat.shape, dtype=dtype)
+    both[0, :, :] = np.real(mat)
+    both[1, :, :] = np.imag(mat)
+    return both
+
+
+@tf_op("RFFT")
+class RFFTOp:
+    # support more dtype
+    supported_dtypes = [
+        onnx_pb.TensorProto.FLOAT,
+        onnx_pb.TensorProto.FLOAT16,
+        onnx_pb.TensorProto.DOUBLE,
+        onnx_pb.TensorProto.COMPLEX64,
+        onnx_pb.TensorProto.COMPLEX128,
+    ]
+
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        """
+        Inspired from `Python implementation of RFFT
+        <https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/>`_.
+
+        Complex version:
+
+        ::
+
+            import numpy as np
+
+            def _DFT_cst(N, fft_length):
+                n = np.arange(N)
+                k = n.reshape((N, 1)).astype(np.float64)
+                M = np.exp(-2j * np.pi * k * n / N)
+                return M[:fft_length // 2 + 1]
+
+            def DFT(x, fft_length=None):
+                if len(x.shape) == 1:
+                    x = x.reshape((-1, 1))
+                else:
+                    x = x.T
+                if fft_length is None:
+                    fft_length = x.shape[0]
+                cst = _DFT_cst(x.shape[0], fft_length)
+                return np.dot(cst, x).T
+
+        Real version, first axis is (real, imag) part:
+
+        ::
+
+            import numpy as np
+
+            def _DFT_real_cst(N, fft_length):
+                n = np.arange(N)
+                k = n.reshape((N, 1)).astype(np.float64)
+                M = np.exp(-2j * np.pi * k * n / N)
+                M = M[:fft_length // 2 + 1]
+                both = np.empty((2,) + M.shape)
+                both[0, :, :] = np.real(M)
+                both[1, :, :] = np.imag(M)
+                return both
+
+            def DFT_real(x, fft_length=None):
+                if len(x.shape) == 1:
+                    x = x.reshape((-1, 1))
+                else:
+                    x = x.T
+                if fft_length is None:
+                    fft_length = x.shape[0]
+                cst = _DFT_real_cst(x.shape[0], fft_length)
+                res = np.dot(cst, x)
+                return np.transpose(res, (0, 2, 1))
+        """
+
+        onnx_dtype = ctx.get_dtype(node.input[0])
+        shape = ctx.get_shape(node.input[0])
+        np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
+        shape_n = shape[-1]
+        utils.make_sure(len(node.input) == 2, "Two inputs expected not %r", len(node.input))
+
+        # This input should be a constant.
+        fft_length_name = node.input[1]
+        node_fft_length = ctx.get_node_by_output(fft_length_name, search_in_parent_graphs=True)
+        utils.make_sure(node_fft_length.type == 'Const',
+                        "fft_length should be a constant, the other case is not implemented yet.")
+        value = node_fft_length.get_attr("value")
+        value_array = to_array(value.t)
+        utils.make_sure(value_array.shape == (1,), "Unexpected shape for fft_length (%r)", value_array.shape)
+        fft_length = value_array[0]
+
+        # TODO: handle this parameter when onnx.helper.make_node is fixed.
+        # Tcomplex = node.get_attr("Tcomplex")
+
+        if np_dtype == np.float16:
+            res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float16)
+            np_dtype = np.float16
+        elif np_dtype in (np.float32, np.complex64):
+            res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float32)
+            np_dtype = np.float32
+        else:
+            res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float64)
+            np_dtype = np.float64
+
+        real_imag_part = make_dft_constant(shape_n, np_dtype, fft_length)
+        onx_real_imag_part = ctx.make_const(
+            name=utils.make_name('cst_rfft_%d' % shape_n), np_val=real_imag_part)
+
+        shapei = list(np.arange(len(shape)))
+        perm = shapei[:-2] + [shapei[-1], shapei[-2]]
+        trx = ctx.make_node(
+            "Transpose", inputs=[node.input[0]], attr=dict(perm=perm),
+            name=utils.make_name(node.name + 'tr'))
+
+        ctx.remove_node(node.name)
+        mult = ctx.make_node(
+            "MatMul", inputs=[onx_real_imag_part.name, trx.output[0]],
+            name=utils.make_name('CPLX_' + node.name + 'rfft'))
+
+        new_shape = [2] + list(shape)
+        shapei = list(np.arange(len(new_shape)))
+        perm = shapei[:-2] + [shapei[-1], shapei[-2]]
+        last_node = ctx.make_node(
+            "Transpose", inputs=[mult.output[0]], attr=dict(perm=perm),
+            name=utils.make_name('CPLX_' + node.name + 'rfft'),
+            shapes=[new_shape], dtypes=[res_onnx_dtype])
+
+        ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()
+
+
+@tf_op("ComplexAbs")
+class ComplexAbsOp:
+    # support more dtype
+    supported_dtypes = [
+        onnx_pb.TensorProto.FLOAT,
+        onnx_pb.TensorProto.FLOAT16,
+        onnx_pb.TensorProto.DOUBLE,
+        onnx_pb.TensorProto.COMPLEX64,
+        onnx_pb.TensorProto.COMPLEX128,
+    ]
+
+    @classmethod
+    def any_version(cls, opset, ctx, node, **kwargs):
+        """
+        Computes the modules of a complex.
+        If the matrix dtype is not complex64 or complex128,
+        it assumes the first dimension means real part (0)
+        and imaginary part (1, :, :...).
+        """
+        onnx_dtype = ctx.get_dtype(node.input[0])
+        shape = ctx.get_shape(node.input[0])
+        np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
+        utils.make_sure(shape[0] == 2, "ComplexAbs expected the first dimension to be 2 but shape is %r", shape)
+
+        ind0 = ctx.make_const(name=utils.make_name('cst0'), np_val=np.array([0], dtype=np.int64))
+        ind1 = ctx.make_const(name=utils.make_name('cst1'), np_val=np.array([1], dtype=np.int64))
+        p2 = ctx.make_const(name=utils.make_name('p2'), np_val=np.array([2], dtype=np_dtype))
+
+        real_part = ctx.make_node(
+            'Gather', inputs=[node.input[0], ind0.name], attr=dict(axis=0),
+            name=utils.make_name('Real_' + node.name))
+        imag_part = ctx.make_node(
+            'Gather', inputs=[node.input[0], ind1.name], attr=dict(axis=0),
+            name=utils.make_name('Imag_' + node.name))
+
+        real_part2 = ctx.make_node(
+            'Pow', inputs=[real_part.output[0], p2.name],
+            name=utils.make_name(real_part.name + 'p2p'))
+
+        imag_part2 = ctx.make_node(
+            'Pow', inputs=[imag_part.output[0], p2.name],
+            name=utils.make_name(imag_part.name + 'p2p'))
+
+        ctx.remove_node(node.name)
+        add = ctx.make_node(
+            "Add", inputs=[real_part2.output[0], imag_part2.output[0]],
+            name=utils.make_name('ComplexAbs_' + node.name))
+
+        if opset == 1:
+            squeezed = ctx.make_node(
+                "Squeeze", inputs=add.output[:1], attr=dict(axes=[0]),
+                name=utils.make_name('ComplexAbs' + node.name))
+        else:
+            squeezed = ctx.make_node(
+                "Squeeze", inputs=[add.output[0], ind0],
+                name=utils.make_name('ComplexAbsSqr' + node.name))
+
+        last_node = ctx.make_node(
+            "Sqrt", inputs=squeezed.output[:1],
+            name=utils.make_name('ComplexAbs' + node.name),
+            shapes=[shape[1:]], dtypes=[onnx_dtype])
+
+        ctx.replace_all_inputs(node.output[0], last_node.output[0])  # ops=ctx.get_nodes()
+
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        cls.any_version(1, ctx, node, **kwargs)
+
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        cls.any_version(11, ctx, node, **kwargs)
diff --git a/tf2onnx/tf_utils.py b/tf2onnx/tf_utils.py
@@ -234,7 +234,11 @@ def tflist_to_onnx(g, shape_override, const_node_values=None):
                     "T_threshold", "element_dtype", "shape_type", "_lower_using_switch_merge",
                     "parallel_iterations", "_num_original_outputs", "output_types", "output_shapes",
                     "key_dtype", "value_dtype", "Tin", "Tout", "capacity", "component_types", "shapes",
-                    "Toutput_types"}
+                    "Toutput_types",
+                    "Tcomplex", "Treal",  # For RFFT, Tcomplex is ignored because
+                                          # onnx.helper.make_node fails,
+                                          # TODO: it should be added back.
+                    }
 
     node_list = g.get_operations()
     functions = {}
diff --git a/tf2onnx/utils.py b/tf2onnx/utils.py
@@ -40,6 +40,8 @@
     onnx_pb.TensorProto.INT64: np.int64,
     onnx_pb.TensorProto.UINT64: np.uint64,
     onnx_pb.TensorProto.BOOL: np.bool,
+    onnx_pb.TensorProto.COMPLEX64: np.complex64,
+    onnx_pb.TensorProto.COMPLEX128: np.complex128,
 }
 
 #
@@ -56,7 +58,9 @@
     onnx_pb.TensorProto.UINT16: "uint16",
     onnx_pb.TensorProto.INT64: "int64",
     onnx_pb.TensorProto.STRING: "string",
-    onnx_pb.TensorProto.BOOL: "bool"
+    onnx_pb.TensorProto.BOOL: "bool",
+    onnx_pb.TensorProto.COMPLEX64: "complex64",
+    onnx_pb.TensorProto.COMPLEX128: "complex128"
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,8 @@`
`40`	`40`	`onnx_pb.TensorProto.INT64: np.int64,`
`41`	`41`	`onnx_pb.TensorProto.UINT64: np.uint64,`
`42`	`42`	`onnx_pb.TensorProto.BOOL: np.bool,`
	`43`	`+ onnx_pb.TensorProto.COMPLEX64: np.complex64,`
	`44`	`+ onnx_pb.TensorProto.COMPLEX128: np.complex128,`
`43`	`45`	`}`
`44`	`46`
`45`	`47`	`#`
`@@ -56,7 +58,9 @@`
`56`	`58`	`onnx_pb.TensorProto.UINT16: "uint16",`
`57`	`59`	`onnx_pb.TensorProto.INT64: "int64",`
`58`	`60`	`onnx_pb.TensorProto.STRING: "string",`
`59`		`- onnx_pb.TensorProto.BOOL: "bool"`
	`61`	`+ onnx_pb.TensorProto.BOOL: "bool",`
	`62`	`+ onnx_pb.TensorProto.COMPLEX64: "complex64",`
	`63`	`+ onnx_pb.TensorProto.COMPLEX128: "complex128"`
`60`	`64`	`}`
`61`	`65`
`62`	`66`