def make_const(self, name, np_val, skip_conversion=False, raw=True): """Make a new constant in the graph. Args: name: const node name, must be unique. np_val: value of type numpy ndarray. skip_conversion: bool, indicate whether this created node would be mapped during conversion. raw: whether to store data at field of raw_data or the specific field according to its dtype """ if raw: onnx_tensor = numpy_helper.from_array(np_val, name) else: onnx_tensor = helper.make_tensor(name, utils.map_numpy_to_onnx_dtype( np_val.dtype), np_val.shape, np_val, raw=False) node = self.make_node("Const", [], outputs=[name], name=name, attr={"value": onnx_tensor}, skip_conversion=skip_conversion) self.set_shape(name, np_val.shape) self.set_dtype(name, utils.map_numpy_to_onnx_dtype(np_val.dtype)) return node
def _replace_node_with_const(node, graph, vals): utils.make_sure( len(node.output) == len(vals), "length of node outputs and const vals should be same") for old_input, val in zip(node.output, vals): const_node = graph.make_const(utils.make_name("const_fold_opt"), val) graph.set_dtype(const_node.output[0], utils.map_numpy_to_onnx_dtype(val.dtype)) graph.set_shape(const_node.output[0], val.shape) graph.replace_all_inputs(graph.get_nodes(), old_input, const_node.output[0]) graph.remove_node(node.name)
def make_const(self, name, np_val, skip_conversion=False): """Make a new constant in the graph. Args: name: const node name, must be unique. np_val: value of type numpy ndarray. skip_conversion: bool, indicate whether this created node would be mapped during conversion. """ onnx_tensor = numpy_helper.from_array(np_val, name) node = self.make_node("Const", [], outputs=[name], name=name, attr={"value": onnx_tensor}, skip_conversion=skip_conversion) self.set_shape(name, np_val.shape) self.set_dtype(name, utils.map_numpy_to_onnx_dtype(np_val.dtype)) return node
def any_version(cls, const_length, opset, 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)) """ supported_dtypes = [ onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.FLOAT16, onnx_pb.TensorProto.DOUBLE, onnx_pb.TensorProto.COMPLEX64, onnx_pb.TensorProto.COMPLEX128, ] consumers = ctx.find_output_consumers(node.output[0]) consumer_types = set(op.type for op in consumers) utils.make_sure( consumer_types == {'ComplexAbs'}, "Current implementation of RFFT or FFT only allows ComplexAbs as consumer not %r", consumer_types) input_name = node.input[0] onnx_dtype = ctx.get_dtype(input_name) utils.make_sure(onnx_dtype in supported_dtypes, "Unsupported input type.") shape = ctx.get_shape(node.input[0]) shape_n = shape[-1] if onnx_dtype in (onnx_pb.TensorProto.COMPLEX64, onnx_pb.TensorProto.COMPLEX128): parent = ctx.get_node_by_output_in_current_graph(node.input[0]) utils.make_sure( parent.type == 'Cast' and parent.get_attr_value('to') == onnx_dtype, "Current implementation of FFT or RFFT assumes the input is real or complex produced " "by a node Cast just before this one.") input_name = parent.input[0] onnx_dtype = ctx.get_dtype(input_name) np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype) 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 if const_length: # RFFT: length of FFT is known, some computation # (see function make_dft_constant) # can be done at conversion time and stored as constant 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") 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) onx_real_imag_part_name = onx_real_imag_part.name else: # FFT: length of FFT is unknown, the matrix # created by function make_dft_constant must be # done in ONNX. dyn_shape_all = ctx.make_node("Shape", inputs=[input_name], name=utils.make_name('CPLX_' + node.name + 'shape')) m1_cst = ctx.make_const(name=utils.make_name('CPLX_m1'), np_val=np.array([-1], dtype=np.int64)) dyn_shape = ctx.make_node('Gather', inputs=[dyn_shape_all.output[0], m1_cst.name]) one_tensor = helper.make_tensor("value", res_onnx_dtype, dims=[1], vals=[1]) cst_1 = ctx.make_node("ConstantOfShape", inputs=[dyn_shape.output[0]], attr={"value": one_tensor}) just_0 = ctx.make_const(name=utils.make_name('CPLX1'), np_val=np.array([0], dtype=np.int64)) rng1 = ctx.make_node("CumSum", inputs=[cst_1.output[0], just_0.name], name=utils.make_name('CPLX_' + node.name + 'range')) p1_cst = ctx.make_const(name=utils.make_name('CPLX_p1'), np_val=np.array([1], dtype=np_dtype)) rng = ctx.make_node("Sub", inputs=[rng1.output[0], p1_cst.name], name=utils.make_name('CPLX_' + node.name + 'range')) resh_cst = ctx.make_const(name=utils.make_name('CPLX_reshape'), np_val=np.array([1, -1], dtype=np.int64)) rng_tr1 = ctx.make_node("Reshape", inputs=[rng.output[0], resh_cst.name], name=utils.make_name('CPLX_' + node.name + 'range')) resh_cst = ctx.make_const(name=utils.make_name('CPLX_reshape'), np_val=np.array([-1, 1], dtype=np.int64)) rng_tr2 = ctx.make_node("Reshape", inputs=[rng.output[0], resh_cst.name], name=utils.make_name('CPLX_' + node.name + 'range')) rng_mat = ctx.make_node('MatMul', inputs=[rng_tr2.output[0], rng_tr1.output[0]], name=utils.make_name('CPLX_' + node.name + 'range2')) pi_cst = ctx.make_const(name=utils.make_name('CPLX_pi'), np_val=np.array([np.pi * 2], dtype=np_dtype)) angle_pi = ctx.make_node("Mul", inputs=[rng_mat.output[0], pi_cst.name], name=utils.make_name('CPLX_' + node.name + 'angle_pi')) shape_cast = ctx.make_node('Cast', inputs=[dyn_shape.output[0]], attr={'to': res_onnx_dtype}) angle_pibn = ctx.make_node("Div", inputs=[angle_pi.output[0], shape_cast.output[0]], name=utils.make_name('CPLX_' + node.name + 'angle')) if opset >= 13: angle = ctx.make_node("Unsqueeze", inputs=[angle_pibn.output[0], just_0.name], name=utils.make_name('CPLX_' + node.name + 'angles')) else: angle = ctx.make_node("Unsqueeze", inputs=[angle_pibn.output[0]], name=utils.make_name('CPLX_' + node.name + 'angles'), attr={'axes': [0]}) rng_cos = ctx.make_node("Cos", inputs=[angle.output[0]], name=utils.make_name('CPLX_' + node.name + 'cos')) rng_sin = ctx.make_node("Sin", inputs=[angle.output[0]], name=utils.make_name('CPLX_' + node.name + 'sin')) onx_real_imag_part = ctx.make_node("Concat", inputs=[rng_cos.output[0], rng_sin.output[0]], name=utils.make_name('CPLX_' + node.name + '_cst_fft'), attr={'axis': 0}) onx_real_imag_part_name = onx_real_imag_part.output[0] shapei = list(np.arange(len(shape))) perm = shapei[:-2] + [shapei[-1], shapei[-2]] trx = ctx.make_node( "Transpose", inputs=[input_name], 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()
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)) """ supported_dtypes = [ onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.FLOAT16, onnx_pb.TensorProto.DOUBLE, onnx_pb.TensorProto.COMPLEX64, onnx_pb.TensorProto.COMPLEX128, ] consumers = ctx.find_output_consumers(node.output[0]) consumer_types = set(op.type for op in consumers) utils.make_sure( consumer_types == {'ComplexAbs'}, "Current implementation of RFFT only allows ComplexAbs as consumer not %r", consumer_types) onnx_dtype = ctx.get_dtype(node.input[0]) utils.make_sure(onnx_dtype in supported_dtypes, "Unsupported input type.") 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()