Update DFT definition to more closely align with PyTorch by enabling axis attribute, and arbitrary tensor rank. (#10842)

* Add axis attribute

* fix breaks

* Enable axis-specified DFT

* remove static cast

Co-authored-by: Sheil Kumar <sheilk@microsoft.com>
This commit is contained in:
Sheil Kumar 2022-03-15 15:27:12 -07:00 committed by GitHub
parent de6d1fcb41
commit 860f28254e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 374 additions and 167 deletions

View file

@ -39,20 +39,16 @@ ONNX_OPERATOR_KERNEL_EX(
kMSExperimentalDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder().MayInplace(0, 0).TypeConstraint("T", BuildKernelDefConstraints<float, double>()),
KernelDefBuilder().MayInplace(0, 0).TypeConstraint("T1", BuildKernelDefConstraints<float, double>())
.TypeConstraint("T2", BuildKernelDefConstraints<int64_t>()),
STFT);
static bool is_real_valued_signal(const onnxruntime::TensorShape & shape) {
// The first dimention is the batch size
// The second dimention is the signal value
return shape.NumDimensions() == 2;
return shape.NumDimensions() == 2 || shape[shape.NumDimensions() - 1] == 1;
}
static bool is_complex_valued_signal(const onnxruntime::TensorShape& shape) {
// The first dimention is the batch size
// The second dimention is the signal length
// The third dimention is set to 2 and represents the real and imaginary parts of the complex sample
return shape.NumDimensions() == 3 && shape[2] == 2;
return shape.NumDimensions() > 2 && shape[shape.NumDimensions() - 1] == 2;
}
static bool is_power_of_2(size_t size) {
@ -143,24 +139,27 @@ static T compute_angular_velocity(size_t number_of_samples, bool inverse) {
}
template <typename T, typename U>
static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx,
const Tensor* X, Tensor* Y, const Tensor* window, bool is_onesided, bool inverse,
static Status fft_radix2(OpKernelContext* /*ctx*/,
const Tensor* X, Tensor* Y,
size_t X_offset, size_t X_stride, size_t Y_offset, size_t Y_stride, int64_t axis,
const Tensor* window, bool is_onesided, bool inverse,
std::vector<std::complex<T>>& V,
std::vector<std::complex<T>>& temp_output) {
// Get shape and significant bits
const auto& X_shape = X->Shape();
size_t number_of_samples = static_cast<size_t>(X_shape[1]);
size_t number_of_samples = static_cast<size_t>(X_shape[axis]);
unsigned significant_bits = static_cast<unsigned>(log2(number_of_samples));
// Get data
auto* X_data = const_cast<U*>(reinterpret_cast<const U*>(X->DataRaw())) + (batch_idx * number_of_samples);
auto* X_data = const_cast<U*>(reinterpret_cast<const U*>(X->DataRaw())) + X_offset;
// Get window
U* window_data = nullptr;
if (window) {
window_data = const_cast<U*>(reinterpret_cast<const U*>(window->DataRaw()));
}
size_t Y_data_stride = 1;
std::complex<T>* Y_data;
if (is_onesided) {
if (temp_output.size() != number_of_samples) {
@ -168,7 +167,8 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx,
}
Y_data = temp_output.data();
} else {
Y_data = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + (batch_idx * number_of_samples);
Y_data = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + Y_offset;
Y_data_stride = Y_stride;
}
auto angular_velocity = compute_angular_velocity<T>(number_of_samples, inverse);
@ -184,9 +184,9 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx,
for (size_t i = 0; i < number_of_samples; i++) {
size_t bit_reversed_index = bit_reverse(i, significant_bits);
auto x = *(X_data + bit_reversed_index);
auto x = *(X_data + bit_reversed_index*X_stride);
auto window_element = window_data ? *(window_data + bit_reversed_index) : 1;
*(Y_data + i) = std::complex<T>(1, 0) * x * window_element;
*(Y_data + i*Y_data_stride) = std::complex<T>(1, 0) * x * window_element;
}
// Run fft_radix2
@ -199,8 +199,8 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx,
auto first_idx = bit_reverse(k, current_significant_bits);
auto second_idx = bit_reverse(midpoint + k, current_significant_bits);
for (size_t j = 0; j < number_of_samples; j += i) {
std::complex<T>* even = (Y_data + j) + k;
std::complex<T>* odd = (Y_data + j) + (midpoint + k);
std::complex<T>* even = (Y_data + j*Y_data_stride) + k;
std::complex<T>* odd = (Y_data + j*Y_data_stride) + (midpoint + k);
std::complex<T> first = *even + (V[first_idx] * *odd);
std::complex<T> second = *even + (V[second_idx] * *odd);
*even = first;
@ -212,32 +212,34 @@ static Status fft_radix2(OpKernelContext* /*ctx*/, size_t batch_idx,
// Scale the output if inverse
if (inverse) {
for (size_t i = 0; i < number_of_samples; i++) {
std::complex<T>& val = *(Y_data + i);
std::complex<T>& val = *(Y_data + i * Y_data_stride);
val /= static_cast<T>(number_of_samples);
}
}
if (is_onesided) {
const auto& Y_shape = Y->Shape();
size_t fft_output_size = static_cast<size_t>(Y_shape[1]);
auto destination = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + (batch_idx * fft_output_size);
memcpy(destination, Y_data, sizeof(std::complex<T>) * fft_output_size);
auto destination = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + Y_offset;
for (size_t i = 0; i < number_of_samples; i++) {
*(destination + Y_stride * i) = *(Y_data + i);
}
}
return Status::OK();
}
template <typename T, typename U>
static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tensor* window, bool inverse) {
static Status dft_naive(const Tensor* X, Tensor* Y,
size_t X_offset, size_t X_stride, size_t Y_offset, size_t Y_stride, int64_t axis,
const Tensor* window, bool inverse) {
// Get shape and significant bits
const auto& X_shape = X->Shape();
size_t number_of_samples = static_cast<size_t>(X_shape[1]);
size_t number_of_samples = static_cast<size_t>(X_shape[axis]);
const auto& Y_shape = Y->Shape();
size_t dft_output_size = static_cast<size_t>(Y_shape[1]);
size_t dft_output_size = static_cast<size_t>(Y_shape[axis]);
// Get data
auto* X_data = const_cast<U*>(reinterpret_cast<const U*>(X->DataRaw())) + (batch_idx * number_of_samples);
auto* Y_data = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + (batch_idx * dft_output_size);
auto* X_data = const_cast<U*>(reinterpret_cast<const U*>(X->DataRaw())) + X_offset;
auto* Y_data = reinterpret_cast<std::complex<T>*>(Y->MutableDataRaw()) + Y_offset;
U* window_data = nullptr;
if (window) {
@ -247,14 +249,14 @@ static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tens
auto angular_velocity = compute_angular_velocity<T>(number_of_samples, inverse);
for (size_t i = 0; i < dft_output_size; i++) {
std::complex<T>& out = *(Y_data + i);
std::complex<T>& out = *(Y_data + i*Y_stride);
out.real(0);
out.imag(0);
for (size_t j = 0; j < number_of_samples; j++) { // vectorize over this loop
auto exponential = std::complex<T>(cos(i * j * angular_velocity), sin(i * j * angular_velocity));
auto window_element = window_data ? * (window_data + j) : 1;
auto element = *(X_data + j) * window_element;
auto element = *(X_data + j*X_stride) * window_element;
out += exponential * element;
}
@ -267,26 +269,65 @@ static Status dft_naive(size_t batch_idx, const Tensor* X, Tensor* Y, const Tens
}
template <typename T, typename U>
static Status discrete_fourier_transform(OpKernelContext* ctx, const Tensor* X, Tensor* Y, const Tensor* window, bool is_onesided, bool inverse,
static Status discrete_fourier_transform(OpKernelContext* ctx, const Tensor* X, Tensor* Y, int64_t axis, const Tensor* window, bool is_onesided, bool inverse,
std::vector<std::complex<T>>& V, std::vector<std::complex<T>>& temp_output) {
// Get shape
const auto& X_shape = X->Shape();
size_t number_of_batches = static_cast<size_t>(X_shape[0]);
size_t number_of_samples = static_cast<size_t>(X_shape[1]);
// radix 2 fft
for (size_t i = 0; i < number_of_batches; i++) {
if (is_power_of_2(number_of_samples)) {
ORT_RETURN_IF_ERROR((fft_radix2<T, U>(ctx, i, X, Y, window, is_onesided, inverse, V, temp_output)));
} else {
ORT_RETURN_IF_ERROR((dft_naive<T, U>(i, X, Y, window, inverse)));
}
const auto& Y_shape = Y->Shape();
size_t number_of_samples = static_cast<size_t>(X_shape[axis]);
auto batch_and_signal_rank = X->Shape().NumDimensions();
auto total_dfts = static_cast<size_t>(X->Shape().Size() / X->Shape()[axis]);
if (X->Shape().NumDimensions() > 2)
{
total_dfts /= X->Shape()[X->Shape().NumDimensions() - 1];
batch_and_signal_rank -= 1;
}
// Calculate x/y offsets/strides
for (size_t i = 0; i < total_dfts; i++)
{
size_t X_offset = 0;
size_t X_stride = X_shape.SizeFromDimension(axis+1);
size_t cumulative_packed_stride = total_dfts;
size_t temp = i;
for (size_t r = 0; r < batch_and_signal_rank; r++) {
if (r == static_cast<size_t>(axis))
{
continue;
}
cumulative_packed_stride /= X_shape[r];
auto index = temp / cumulative_packed_stride;
temp -= (index * cumulative_packed_stride);
X_offset += index * X_shape.SizeFromDimension(r + 1);
}
size_t Y_offset = 0;
size_t Y_stride = Y_shape.SizeFromDimension(axis + 1) / 2;
cumulative_packed_stride = total_dfts;
temp = i;
for (size_t r = 0; r < batch_and_signal_rank; r++) {
if (r == static_cast<size_t>(axis))
{
continue;
}
cumulative_packed_stride /= X_shape[r];
auto index = temp / cumulative_packed_stride;
temp -= (index * cumulative_packed_stride);
Y_offset += index * Y_shape.SizeFromDimension(r + 1) / 2;
}
if (is_power_of_2(number_of_samples)) {
ORT_RETURN_IF_ERROR((fft_radix2<T, U>(ctx, X, Y, X_offset, X_stride, Y_offset, Y_stride, axis, window, is_onesided, inverse, V, temp_output)));
} else {
ORT_RETURN_IF_ERROR((dft_naive<T, U>(X, Y, X_offset, X_stride, Y_offset, Y_stride, axis, window, inverse)));
}
}
return Status::OK();
}
static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided, bool inverse) {
static Status discrete_fourier_transform(OpKernelContext* ctx, int64_t axis, bool is_onesided, bool inverse) {
// Get input shape
const auto* X = ctx->Input<Tensor>(0);
const auto& X_shape = X->Shape();
@ -295,13 +336,21 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided,
// Get the DFT output size. Onesided will return only the unique values!
// note: x >> 1 === std::floor(x / 2.f)
int64_t number_of_samples = static_cast<int64_t>(X_shape[1]);
int64_t number_of_samples = static_cast<int64_t>(X_shape[axis]);
auto dft_output_size = is_onesided ?
((number_of_samples >> 1) + 1) :
number_of_samples;
// Get output shape
auto Y_shape = onnxruntime::TensorShape({X_shape[0], dft_output_size, 2});
auto Y_shape = onnxruntime::TensorShape(X_shape);
if (X_shape.NumDimensions() == 2)
{
Y_shape = onnxruntime::TensorShape({X_shape[0], dft_output_size, 2});
} else
{
Y_shape[Y_shape.NumDimensions() - 1] = 2;
}
Y_shape[axis] = dft_output_size;
auto Y = ctx->Output(0, Y_shape);
// Get data type
@ -312,9 +361,9 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided,
std::vector<std::complex<float>> V;
std::vector<std::complex<float>> temp_output;
if (is_real_valued) {
ORT_RETURN_IF_ERROR((discrete_fourier_transform<float, float>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output)));
ORT_RETURN_IF_ERROR((discrete_fourier_transform<float, float>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output)));
} else if (is_complex_valued) {
ORT_RETURN_IF_ERROR((discrete_fourier_transform<float, std::complex<float>>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output)));
ORT_RETURN_IF_ERROR((discrete_fourier_transform<float, std::complex<float>>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output)));
} else {
ORT_THROW("Unsupported input signal shape. The signal's first dimenstion must be the batch dimension and its second dimension must be the signal length dimension. It may optionally include a 3rd dimension of size 2 for complex inputs.", data_type);
}
@ -322,9 +371,9 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided,
std::vector<std::complex<double>> V;
std::vector<std::complex<double>> temp_output;
if (is_real_valued) {
ORT_RETURN_IF_ERROR((discrete_fourier_transform<double, double>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output)));
ORT_RETURN_IF_ERROR((discrete_fourier_transform<double, double>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output)));
} else if (is_complex_valued) {
ORT_RETURN_IF_ERROR((discrete_fourier_transform<double, std::complex<double>>(ctx, X, Y, nullptr, is_onesided, inverse, V, temp_output)));
ORT_RETURN_IF_ERROR((discrete_fourier_transform<double, std::complex<double>>(ctx, X, Y, axis, nullptr, is_onesided, inverse, V, temp_output)));
} else {
ORT_THROW("Unsupported input signal shape. The signal's first dimenstion must be the batch dimension and its second dimension must be the signal length dimension. It may optionally include a 3rd dimension of size 2 for complex inputs.", data_type);
}
@ -336,12 +385,12 @@ static Status discrete_fourier_transform(OpKernelContext* ctx, bool is_onesided,
}
Status DFT::Compute(OpKernelContext* ctx) const {
ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, is_onesided_, false));
ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, axis_ + 1, is_onesided_, false));
return Status::OK();
}
Status IDFT::Compute(OpKernelContext* ctx) const {
ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, false, true));
ORT_RETURN_IF_ERROR(discrete_fourier_transform(ctx, axis_ + 1, false, true));
return Status::OK();
}
@ -376,9 +425,9 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside
// Get signal
const auto* signal = ctx->Input<Tensor>(0);
const auto* window = ctx->Input<Tensor>(1);
const auto* frame_length_tensor = ctx->Input<Tensor>(2);
const auto frame_step = get_scalar_value_from_tensor<int64_t>(ctx->Input<Tensor>(3));
const auto frame_step = get_scalar_value_from_tensor<int64_t>(ctx->Input<Tensor>(1));
const auto* window = ctx->Input<Tensor>(2);
const auto* frame_length_tensor = ctx->Input<Tensor>(3);
// Get input signal shape
const auto& signal_shape = signal->Shape();
@ -468,7 +517,7 @@ static Status short_time_fourier_transform(OpKernelContext* ctx, bool is_oneside
0);
// Run individual dft
ORT_RETURN_IF_ERROR((discrete_fourier_transform<T, U>(ctx, &input, &output, window, is_onesided, false, V, temp_output)));
ORT_RETURN_IF_ERROR((discrete_fourier_transform<T, U>(ctx, &input, &output, 1, window, is_onesided, false, V, temp_output)));
}
}

View file

@ -8,16 +8,20 @@ namespace contrib {
class DFT final : public OpKernel {
bool is_onesided_ = true;
int64_t axis_ = 0;
public:
explicit DFT(const OpKernelInfo& info) : OpKernel(info) {
is_onesided_ = static_cast<bool>(info.GetAttrOrDefault<int64_t>("onesided", 0));
axis_ = info.GetAttrOrDefault<int64_t>("axis", 0);
}
Status Compute(OpKernelContext* ctx) const override;
};
class IDFT final : public OpKernel {
int64_t axis_ = 0;
public:
explicit IDFT(const OpKernelInfo& info) : OpKernel(info) {
axis_ = info.GetAttrOrDefault<int64_t>("axis", 0);
}
Status Compute(OpKernelContext* ctx) const override;
};

View file

@ -42,6 +42,24 @@ static T get_scalar_value_from_tensor(const ONNX_NAMESPACE::TensorProto* t) {
}
}
inline const ONNX_NAMESPACE::TensorShapeProto* getOptionalInputShape(ONNX_NAMESPACE::InferenceContext& ctx, size_t n) {
const auto* input_type = ctx.getInputType(n);
if (input_type == nullptr) {
return nullptr;
}
const auto value_case = input_type->value_case();
if (value_case != ONNX_NAMESPACE::TypeProto::kTensorType && value_case != ONNX_NAMESPACE::TypeProto::kSparseTensorType) {
fail_type_inference("Attribute expected to have tensor or sparse tensor type");
}
if (value_case == ONNX_NAMESPACE::TypeProto::kTensorType) {
return &input_type->tensor_type().shape();
} else {
return &input_type->sparse_tensor_type().shape();
}
}
void RegisterSignalSchemas() {
MS_SIGNAL_OPERATOR_SCHEMA(DFT)
.SetDomain(kMSExperimentalDomain)
@ -53,132 +71,242 @@ void RegisterSignalSchemas() {
"Values can be 0 or 1.",
AttributeProto::AttributeType::AttributeProto_AttributeType_INT,
static_cast<int64_t>(0))
.Attr("axis",
"The axis on which to perform the DFT. By default this value is set to 0, which corresponds to the first dimension after the batch index."
"This value must be less than signal_dimN, where signal_dimN is the number of dimensions in the signal.",
AttributeProto::AttributeType::AttributeProto_AttributeType_INT,
static_cast<int64_t>(0))
.Input(0,
"input",
"For complex input, the following shape is expected: [batch_idx][n_fft][2]"
"The final dimension represents the real and imaginary parts of the value."
"For real input, the following shape is expected: [batch_idx][n_fft]"
"Thefirstdimensionisthebatchdimension.",
"T")
"input",
"For real input, the following shape is expected: [batch_idx][n_fft]."
"For complex input, the following shape is expected: [batch_idx][n_fft][2]."
"The final dimension represents the real and imaginary parts of the value."
"For real multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][1]."
"For complex multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][2]."
"The first dimension is the batch dimension.",
"T")
.Output(0,
"output",
"The Fourier Transform of the input vector."
"If onesided is 1, [batch_idx][floor(n_fft/2)+1][2]"
"If onesided is 0, [batch_idx][n_fft][2]",
"If signal_dimN = 1, and onesided is 0, [batch_idx][n_fft][2]"
"If signal_dimN = 1, and onesided is 1, [batch_idx][floor(n_fft/2)+1][2]"
"If signal_dimN = 2, and onesided is 0 and axis = 0, [batch_idx][signal_dim1][signal_dim2][2]"
"If signal_dimN = 2, and onesided is 0 and axis = 1, [batch_idx][signal_dim1][signal_dim2][2]"
"If signal_dimN = 2, and onesided is 1 and axis = 0, [batch_idx][floor(signal_dim1/2)+1][signal_dim2][2]"
"If signal_dimN = 2, and onesided is 1 and axis = 1, [batch_idx][signal_dim1][floor(signal_dim2/2)+1][2]",
"T")
.TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "")
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input and output types to float tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
int64_t ndim = 1;
propagateElemTypeFromInputToOutput(ctx, 0, 0);
const int64_t batch_ndim = 1;
bool is_onesided = true;
auto attr_proto = ctx.getAttribute("onesided");
if (attr_proto && attr_proto->has_i()) {
is_onesided = static_cast<bool>(attr_proto->i());
}
if (ctx.getInputType(0)->tensor_type().has_shape()) {
auto& input_shape = getInputShape(ctx, 0);
ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape;
if (is_onesided) {
auto n_fft = input_shape.dim(1).dim_value();
result_shape.mutable_dim(1)->set_dim_value((n_fft >> 1) + 1);
}
auto dim_size = static_cast<int64_t>(input_shape.dim_size());
if (dim_size == ndim + 1) { // real input
result_shape.add_dim()->set_dim_value(2); // output is same shape, but with extra dim for 2 values (real/imaginary)
} else if (dim_size == ndim + 2) { // complex input, do nothing
} else {
fail_shape_inference(
"the input_shape must [batch_idx][n_fft] for real values or [batch_idx][n_fft][2] for complex values.")
auto has_component_dimension = dim_size > 2;
ONNX_NAMESPACE::TensorShapeProto result_shape_proto = input_shape;
bool is_onesided = static_cast<bool>(getAttribute(ctx, "onesided", 0));
if (is_onesided) {
// Since signal_ndim = 1, and multidimensional DFT is not supported,
// only the single signal dim (1) needs to be updated
auto n_fft = input_shape.dim(1).dim_value();
result_shape_proto.mutable_dim(1)->set_dim_value((n_fft >> 1) + 1);
}
updateOutputShape(ctx, 0, result_shape);
}
if (has_component_dimension) {
result_shape_proto.mutable_dim(static_cast<int>(dim_size - 1))->set_dim_value(2);
} else {
result_shape_proto.add_dim()->set_dim_value(2);
}
updateOutputShape(ctx, 0, result_shape_proto);
});
;
MS_SIGNAL_OPERATOR_SCHEMA(IDFT)
.SetDomain(kMSExperimentalDomain)
.SinceVersion(1)
.SetDoc(R"DOC(IDFT)DOC")
.Attr("axis",
"The axis on which to perform the DFT. By default this value is set to 0, which corresponds to the first dimension after the batch index."
"This value must be less than signal_dimN, where signal_dimN is the number of dimensions in the signal.",
AttributeProto::AttributeType::AttributeProto_AttributeType_INT,
static_cast<int64_t>(0))
.Input(0,
"input",
"A complex signal of dimension signal_ndim."
"The last dimension of the tensor should be 2,"
"representing the real and imaginary components of complex numbers,"
"and should have at least signal_ndim + 2 dimensions."
"For real input, the following shape is expected: [batch_idx][n_fft]."
"For complex input, the following shape is expected: [batch_idx][n_fft][2]."
"The final dimension represents the real and imaginary parts of the value."
"For real multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][1]."
"For complex multi-dimensional input, the following shape is expected: [batch_idx][signal_dim1][signal_dim2]...[signal_dimN][2]."
"The first dimension is the batch dimension.",
"T")
.Output(0,
"output",
"The inverse fourier transform of the input vector,"
"using the same format as the input.",
"T")
.TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "")
"The inverse discrete Fourier transform of the input. "
"If signal_dimN = 1, [batch_idx][n_fft][2]"
"If signal_dimN = 2 and axis = 0, [batch_idx][signal_dim1][signal_dim2][2]"
"If signal_dimN = 2 and axis = 1, [batch_idx][signal_dim1][signal_dim2][2]"
"For all types of input, the last dimension of the output represents the components of a complex number.",
"T",
OpSchema::Single,
true,
1,
OpSchema::NonDifferentiable)
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input and output types to float tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
int64_t ndim = 1;
auto attr_proto = ctx.getAttribute("signal_ndim");
if (attr_proto && attr_proto->has_i()) {
ndim = static_cast<size_t>(attr_proto->i());
}
propagateElemTypeFromInputToOutput(ctx, 0, 0);
const int64_t batch_ndim = 1;
auto& input_shape = getInputShape(ctx, 0);
ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape;
auto dim_size = static_cast<int64_t>(input_shape.dim_size());
auto has_component_dimension = dim_size > 2;
auto& input_shape = getInputShape(ctx, 0);
ONNX_NAMESPACE::TensorShapeProto result_shape = input_shape;
if (has_component_dimension) {
result_shape.mutable_dim(static_cast<int>(dim_size - 1))->set_dim_value(2);
} else {
result_shape.add_dim()->set_dim_value(2);
}
auto dim_size = static_cast<int64_t>(input_shape.dim_size());
if (dim_size == ndim + 1) { // real input
result_shape.add_dim()->set_dim_value(2); // output is same shape, but with extra dim for 2 values (real/imaginary)
} else if (dim_size == ndim + 2) { // complex input, do nothing
} else {
fail_shape_inference(
"the input_shape must have 1 + signal_ndim dimensions for real inputs, or 2 + signal_ndim dimensions for complex input.")
}
updateOutputShape(ctx, 0, result_shape);
updateOutputShape(ctx, 0, result_shape);
});
MS_SIGNAL_OPERATOR_SCHEMA(STFT)
.SetDomain(kMSExperimentalDomain)
.SinceVersion(1)
.SetDoc(R"DOC(STFT)DOC")
.Attr("onesided",
"If True (default), only values for half of the fft size are returned because the real-to-complex Fourier transform satisfies the conjugate symmetry."
"The output tensor will return the first floor(n_fft/2) + 1 values from the DFT."
"Values can be 0 or 1.",
AttributeProto::AttributeType::AttributeProto_AttributeType_INT,
static_cast<int64_t>(1))
.Attr(
"onesided",
"If onesided is 1, only values for w in [0, 1, 2, ..., floor(n_fft/2) + 1] are returned because "
"the real-to-complex Fourier transform satisfies the conjugate symmetry, i.e., X[m, w] = X[m,w]=X[m,n_fft-w]*. "
"Note if the input or window tensors are complex, then onesided output is not possible. "
"Enabling onesided with real inputs performs a Real-valued fast Fourier transform (RFFT)."
"When invoked with real or complex valued input, the default value is 0. "
"Values can be 0 or 1.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0,
"signal",
"A complex signal of dimension signal_ndim."
"The last dimension of the tensor should be 2,"
"representing the real and imaginary components of complex numbers,"
"and should have at least signal_ndim + 2 dimensions."
"The first dimension is the batch dimension.",
"T1")
.Input(1,
"window",
"A tensor representing the window that will be slid over the input signal.",
"Input tensor representing a real or complex valued signal. "
"For real input, the following shape is expected: [batch_size][signal_length]. "
"For complex input, the following shape is expected: [batch_size][signal_length][2], where "
"[batch_size][signal_length][0] represents the real component and [batch_size][signal_length][1] represents the imaginary component of the signal.",
"T1",
OpSchema::FormalParameterOption::Optional)
.Input(2,
"frame_length", // frame_length, fft_length, pad_mode
"Size of the fft.",
"T2",
OpSchema::FormalParameterOption::Optional)
.Input(3,
OpSchema::Single,
true,
1,
OpSchema::NonDifferentiable)
.Input(1,
"frame_step",
"The number of samples to step between successive DFTs.",
"T2")
"T2",
OpSchema::Single,
true,
1,
OpSchema::NonDifferentiable)
.Input(2,
"window",
"A tensor representing the window that will be slid over the signal."
"The window must have rank 1 with shape: [window_shape]. "
"It's an optional value. ",
"T1",
OpSchema::Optional,
true,
1,
OpSchema::NonDifferentiable)
.Input(3,
"frame_length",
"A scalar representing the size of the DFT. "
"It's an optional value.",
"T2",
OpSchema::Optional,
true,
1,
OpSchema::NonDifferentiable)
.Output(0,
"output",
"The inverse fourier transform of the input vector,"
"using the same format as the input.",
"T1")
.TypeConstraint("T1", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "")
.TypeConstraint("T2", {"tensor(int64)"}, "");
.TypeConstraint(
"T1",
{"tensor(float)",
"tensor(float16)",
"tensor(double)",
"tensor(bfloat16)"},
"Constrain signal and output to float tensors.")
.TypeConstraint(
"T2",
{"tensor(int64)"},
"Constrain scalar length types to int64_t.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
constexpr int64_t batch_ndim = 1;
constexpr int64_t component_ndim = 1;
// Get inputs
auto& input_shape = getInputShape(ctx, 0);
auto frame_step = get_scalar_value_from_tensor<int64_t>(ctx.getInputData(1));
const ONNX_NAMESPACE::TensorShapeProto* window_input = nullptr;
try {
window_input = getOptionalInputShape(ctx, 2);
} catch (...) {
window_input = nullptr;
}
const ONNX_NAMESPACE::TensorShapeProto* frame_length_input = nullptr;
try {
frame_length_input = getOptionalInputShape(ctx, 3);
} catch (...) {
frame_length_input = nullptr;
}
// Determine the size of the DFT based on the 2 optional inputs window and frame_length. One must be set.
int64_t dft_size = 0;
if (window_input == nullptr && frame_length_input == nullptr) {
fail_type_inference("STFT expects to have at least one of these inputs set: [window, frame_length].");
} else if (window_input != nullptr && window_input->dim_size() > 0 && frame_length_input != nullptr) {
if (window_input->dim_size() != 1) {
fail_type_inference("STFT's window input, must have rank = 1.");
}
auto window_length = window_input->dim(0).dim_value();
auto frame_length = get_scalar_value_from_tensor<int64_t>(ctx.getInputData(3));
if (window_length != frame_length) {
fail_type_inference("If STFT has both a window input and frame_length specified, the dimension of the window must match the frame_length specified!");
}
dft_size = window_length;
} else if (window_input != nullptr && window_input->dim_size() > 0) {
if (window_input->dim_size() != 1) {
fail_type_inference("STFT's window input, must have rank = 1.");
}
dft_size = window_input->dim(0).dim_value();
} else if (frame_length_input != nullptr) {
dft_size = get_scalar_value_from_tensor<int64_t>(ctx.getInputData(3));
}
bool is_onesided = static_cast<bool>(getAttribute(ctx, "onesided", 0));
if (is_onesided) {
dft_size = is_onesided ? ((dft_size >> 1) + 1) : dft_size;
}
auto signal_size = input_shape.dim(1).dim_value();
auto n_dfts = static_cast<int64_t>(std::floor((signal_size - dft_size) / static_cast<float>(frame_step)) + 1);
// The output has the following shape: [batch_size][frames][dft_unique_bins][2]
ONNX_NAMESPACE::TensorShapeProto result_shape_proto;
result_shape_proto.add_dim()->set_dim_value(input_shape.dim(0).dim_value()); // batch size
result_shape_proto.add_dim()->set_dim_value(n_dfts);
result_shape_proto.add_dim()->set_dim_value(dft_size);
result_shape_proto.add_dim()->set_dim_value(2);
updateOutputShape(ctx, 0, result_shape_proto);
});
// Window Functions
MS_SIGNAL_OPERATOR_SCHEMA(HannWindow)

View file

@ -476,10 +476,13 @@ static void WindowFunction(const wchar_t* window_operator_name, TensorKind kind)
#endif
#if !defined(BUILD_INBOX) && defined(BUILD_MS_EXPERIMENTAL_OPS)
static void DiscreteFourierTransform(bool is_onesided = false) {
std::vector<int64_t> shape = {1, 5};
std::vector<int64_t> output_shape = {1, 5, 2};
output_shape[1] = is_onesided ? (1 + (shape[1] >> 1)) : shape[1];
static void DiscreteFourierTransform(size_t axis, bool is_onesided = false) {
auto axis_dim = axis + 1;
printf("\nDiscrete Fourier Transform [axis=%d, is_onesided=%s]\n", static_cast<int>(axis_dim), is_onesided ? "true" : "false");
std::vector<int64_t> shape = {2, 5, 8, 1};
std::vector<int64_t> output_shape = {2, 5, 8, 2};
output_shape[axis_dim] = is_onesided ? (1 + (shape[axis_dim] >> 1)) : shape[axis_dim];
auto model =
LearningModelBuilder::Create(13)
@ -487,6 +490,7 @@ static void DiscreteFourierTransform(bool is_onesided = false) {
.Outputs().Add(LearningModelBuilder::CreateTensorFeatureDescriptor(L"Output.Spectra", TensorKind::Float, output_shape))
.Operators().Add(Operator(L"DFT", MS_EXPERIMENTAL_DOMAIN)
.SetInput(L"input", L"Input.Signal")
.SetAttribute(L"axis", TensorInt64Bit::CreateFromArray({}, {INT64(axis)}))
.SetAttribute(L"onesided", TensorInt64Bit::CreateFromArray({}, {is_onesided}))
.SetOutput(L"output", L"Output.Spectra"))
.CreateModel();
@ -495,19 +499,38 @@ static void DiscreteFourierTransform(bool is_onesided = false) {
LearningModelBinding binding(session);
// Populate binding
binding.Bind(L"Input.Signal", TensorFloat::CreateFromArray(shape, {1, 2, 3, 4, 5}));
binding.Bind(
L"Input.Signal",
TensorFloat::CreateFromArray(
shape,
{1, 2, 3, 4, 5, 6, 7, 8,
1, 2, 3, 4, 5, 6, 7, 8,
1, 2, 3, 4, 5, 6, 7, 8,
1, 2, 3, 4, 5, 6, 7, 8,
1, 2, 3, 4, 5, 6, 7, 8,
2, 4, 6, 8, 10, 12, 14, 16,
2, 4, 6, 8, 10, 12, 14, 16,
2, 4, 6, 8, 10, 12, 14, 16,
2, 4, 6, 8, 10, 12, 14, 16,
2, 4, 6, 8, 10, 12, 14, 16,
}));
// Evaluate
auto result = session.Evaluate(binding, L"");
// Check results
printf("Output.Spectra\n");
auto y_tensor = result.Outputs().Lookup(L"Output.Spectra").as<TensorFloat>();
auto y_ivv = y_tensor.GetAsVectorView();
for (int i = 0; i < output_shape[0] * output_shape[1] * 2; i += 2) {
printf("(%f + %fi), ", y_ivv.GetAt(i), y_ivv.GetAt(i + 1));
}
printf("\n");
// // Check results
// printf("Output.Spectra\n");
// auto y_tensor = result.Outputs().Lookup(L"Output.Spectra").as<TensorFloat>();
// auto y_ivv = y_tensor.GetAsVectorView();
// for (uint32_t i = 0; i < y_ivv.Size(); i+=2) {
// auto format_size = 16 * (!is_onesided || axis == 0) + 10 * (is_onesided && axis == 1);
// if (i % format_size == 0 && i != 0) {
// printf("\n");
// }
// printf("(%.2f + %.2fi), ", y_ivv.GetAt(i), y_ivv.GetAt(i + 1));
// }
// printf("\n");
}
#endif
@ -612,20 +635,20 @@ static void STFT(size_t batch_size, size_t signal_size, size_t dft_size,
printf("%f, ", window_ivv.GetAt(i));
}
printf("\n");
printf("Output.STFT\n");
// Check results
auto y_tensor = result.Outputs().Lookup(L"Output.STFT").as<TensorFloat>();
auto y_ivv = y_tensor.GetAsVectorView();
auto size = y_ivv.Size();
WINML_EXPECT_EQUAL(size, n_dfts * output_shape[2] * 2);
for (size_t dft_idx = 0; dft_idx < n_dfts; dft_idx++) {
for (size_t i = 0; INT64(i) < output_shape[2]; i++) {
auto real_idx = static_cast<uint32_t>((i * 2) + (2 * dft_idx * output_shape[2]));
printf("(%d, %f , %fi), ", static_cast<uint32_t>(i), y_ivv.GetAt(real_idx), y_ivv.GetAt(real_idx + 1));
}
}
printf("\n");
//printf("Output.STFT\n");
//// Check results
//auto y_tensor = result.Outputs().Lookup(L"Output.STFT").as<TensorFloat>();
//auto y_ivv = y_tensor.GetAsVectorView();
//auto size = y_ivv.Size();
//WINML_EXPECT_EQUAL(size, n_dfts * output_shape[2] * 2);
//for (size_t dft_idx = 0; dft_idx < n_dfts; dft_idx++) {
// for (size_t i = 0; INT64(i) < output_shape[2]; i++) {
// auto real_idx = static_cast<uint32_t>((i * 2) + (2 * dft_idx * output_shape[2]));
// printf("(%d, %f , %fi), ", static_cast<uint32_t>(i), y_ivv.GetAt(real_idx), y_ivv.GetAt(real_idx + 1));
// }
//}
//
//printf("\n");
}
#endif
@ -913,8 +936,11 @@ static void ModelBuilding_ConstantMatmul() {
static void ModelBuilding_DiscreteFourierTransform() {
#if !defined(BUILD_INBOX) && defined(BUILD_MS_EXPERIMENTAL_OPS)
DiscreteFourierTransform(false /*onesided*/);
DiscreteFourierTransform(true /*onesided*/);
DiscreteFourierTransform(0, false /*onesided*/);
DiscreteFourierTransform(0, true /*onesided*/);
DiscreteFourierTransform(1, false /*onesided*/);
DiscreteFourierTransform(1, true /*onesided*/);
#endif
}