Support Dims attribute for FFT (#33)

Author: dkurt

.github/workflows/main.yml

@@ -174,4 +174,4 @@ jobs:
         run: |
           python3 -m pip install --upgrade pip
           python3 -m pip install twine
-          python3 -m twine upload --repository testpypi wheel*/*.whl --skip-existing
+          python3 -m twine upload wheel*/*.whl --skip-existing

examples/fft/export_model.py

@@ -7,31 +7,30 @@
 
 
 class MyModel(nn.Module):
-    def __init__(self):
+    def __init__(self, inverse, centred, dims):
         super(MyModel, self).__init__()
+        self.inverse = inverse
+        self.centred = centred
+        self.dims = dims
         self.fft = FFT()
 
     def forward(self, x):
-        y = self.fft.apply(x, False)
-        y = y * 2
-        # TODO: there is a bug with "inverse" data attribute in OpenVINO 2021.4
-        y = self.fft.apply(y, True)
-        return y
+        return self.fft.apply(x, self.inverse, self.centred, self.dims)
 
 
-def export(shape):
+def export(shape, inverse, centered, dims):
     np.random.seed(324)
     torch.manual_seed(32)
 
-    model = MyModel()
+    model = MyModel(inverse, centered, dims)
     inp = Variable(torch.randn(shape))
     model.eval()
 
     with torch.no_grad():
         torch.onnx.export(model, inp, 'model.onnx',
                           input_names=['input'],
                           output_names=['output'],
-                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH)
 
     ref = model(inp)
     np.save('inp', inp.detach().numpy())

examples/fft/export_model_with_roll.py

@@ -34,39 +34,41 @@ def roll(
     right_part = data.narrow(dim_index, data.size(dims) - shift, shift)
     return torch.cat([right_part, left_part], dim=dim_index)
 
-def fftshift(data: torch.Tensor) -> torch.Tensor:
-    dim = (1, 2)
-    shift = [data.size(curr_dim) // 2 for curr_dim in dim]
-    return roll(data, shift, dim)
+def fftshift(data: torch.Tensor, dims) -> torch.Tensor:
+    shift = [data.size(curr_dim) // 2 for curr_dim in dims]
+    return roll(data, shift, dims)
 
-def ifftshift(data: torch.Tensor) -> torch.Tensor:
-    dim = (1, 2)
-    shift = [(data.size(curr_dim) + 1) // 2 for curr_dim in dim]
-    return roll(data, shift, dim)
+def ifftshift(data: torch.Tensor, dims) -> torch.Tensor:
+    shift = [(data.size(curr_dim) + 1) // 2 for curr_dim in dims]
+    return roll(data, shift, dims)
 
 class FFT(torch.autograd.Function):
     @staticmethod
-    def symbolic(g, x, inverse, centered=False):
-        return g.op('IFFT' if inverse else 'FFT', x,
+    def symbolic(g, x, inverse, centered, dims):
+        dims = torch.tensor(dims)
+        dims = g.op("Constant", value_t=dims)
+
+        return g.op('IFFT' if inverse else 'FFT', x, dims,
                     inverse_i=inverse, centered_i=centered)
 
     @staticmethod
-    def forward(self, x, inverse, centered=False):
+    def forward(self, x, inverse, centered, dims):
         # https://pytorch.org/docs/stable/torch.html#torch.fft
-        signal_ndim = 2 if len(x.shape) == 5 else 1
         if centered:
-                x = ifftshift(x)
+            x = ifftshift(x, dims)
 
         if version.parse(torch.__version__) >= version.parse("1.8.0"):
             func = torch.fft.ifftn if inverse else torch.fft.fftn
             x = torch.view_as_complex(x)
-            y = func(x, dim=list(range(1, signal_ndim + 1)), norm="ortho")
+            y = func(x, dim=dims, norm="ortho")
             y = torch.view_as_real(y)
         else:
+            signal_ndim = max(dims)
+            assert dims == list(range(1, signal_ndim + 1))
             func = torch.ifft if inverse else torch.fft
             y = func(input=x, signal_ndim=signal_ndim, normalized=True)
 
         if centered:
-            y = fftshift(y)
+            y = fftshift(y, dims)
 
         return y

examples/fft/fft.py

@@ -11,6 +11,7 @@ class FFTFrontExtractor(FrontExtractorOp):
     def extract(cls, node):
         data = {
             "inverse": onnx_attr(node, "inverse", "i"),
+            "centered": onnx_attr(node, "centered", "i"),
         }
 
         FFT.update_node_stat(node, data)
@@ -25,6 +26,7 @@ class IFFTFrontExtractor(FrontExtractorOp):
     def extract(cls, node):
         data = {
             "inverse": 1,
+            "centered": onnx_attr(node, "centered", "i"),
         }
 
         IFFT.update_node_stat(node, data)

mo_extensions/front/onnx/fft_ext.py

@@ -13,15 +13,15 @@ def __init__(self, graph: Graph, attrs: dict):
             {
                 "type": __class__.op,
                 "op": __class__.op,
-                "in_ports_count": 1,
+                "in_ports_count": 2,
                 "out_ports_count": 1,
                 "infer": copy_shape_infer,
             },
             attrs,
         )
 
     def supported_attrs(self):
-        return ["inverse"]
+        return ["inverse", "centered"]
 
 
 class IFFT(Op):
@@ -34,12 +34,12 @@ def __init__(self, graph: Graph, attrs: dict):
             {
                 "type": __class__.op,
                 "op": __class__.op,
-                "in_ports_count": 1,
+                "in_ports_count": 2,
                 "out_ports_count": 1,
                 "infer": copy_shape_infer,
             },
             attrs,
         )
 
     def supported_attrs(self):
-        return ["inverse"]
+        return ["inverse", "centered"]

mo_extensions/ops/FFT.py

@@ -60,18 +60,21 @@ def test_unpool_reshape():
     export(mode='dynamic_size', shape=[4, 3, 17, 8])
     run_test(convert_ir=False)
 
-@pytest.mark.parametrize("shape", [[5, 120, 2], [4, 240, 320, 2], [3, 5, 240, 320, 2]])
-def test_fft(shape):
+@pytest.mark.parametrize("shape", [[5, 120, 2], [4, 240, 320, 2], [3, 16, 240, 320, 2]])
+@pytest.mark.parametrize("inverse", [False, True])
+@pytest.mark.parametrize("centered", [False, True])
+@pytest.mark.parametrize("test_onnx", [False, True])
+@pytest.mark.parametrize("dims", [[1], [1, 2], [2, 3]])
+def test_fft(shape, inverse, centered, test_onnx, dims):
     from examples.fft.export_model import export
 
-    export(shape=shape)
-    run_test()
-
-@pytest.mark.parametrize("test_onnx", [True, False])
-def test_fft_roll(test_onnx):
-    from examples.fft.export_model_with_roll import export
+    if len(shape) == 3 and dims != [1] or \
+       len(shape) == 4 and dims == [2, 3] or \
+       len(shape) == 5 and dims == [1] or \
+       centered and len(dims) != 2:
+        pytest.skip("unsupported configuration")
 
-    export()
+    export(shape, inverse, centered, dims)
     run_test(test_onnx=test_onnx)
 
 

tests/run_tests.py

@@ -38,7 +38,7 @@ class FFTImpl : public InferenceEngine::ILayerExecImpl {
                                         std::vector<InferenceEngine::Blob::Ptr> &outputs,
                                         InferenceEngine::ResponseDesc *resp) noexcept override;
 private:
-    ngraph::Shape inpShape;
+    std::vector<ngraph::Shape> inShapes;
     ngraph::Shape outShape;
     bool inverse, centered;
     std::string error;

user_ie_extensions/cpu_kernel.hpp

@@ -21,13 +21,13 @@ Extension::Extension() {
         ngraph::OutputVector ng_inputs {node.get_ng_inputs()};
         bool inverse = node.get_attribute_value<int64_t>("inverse");
         bool centered = node.get_attribute_value<int64_t>("centered");
-        return {std::make_shared<FFTOp>(ng_inputs.at(0), inverse, centered)};
+        return {std::make_shared<FFTOp>(ng_inputs.at(0), ng_inputs.at(1), inverse, centered)};
     });
     ngraph::onnx_import::register_operator(IFFTOp::type_info.name, 1, "", [](const ngraph::onnx_import::Node& node) -> ngraph::OutputVector {
         ngraph::OutputVector ng_inputs {node.get_ng_inputs()};
         bool inverse = node.get_attribute_value<int64_t>("inverse");
         bool centered = node.get_attribute_value<int64_t>("centered");
-        return {std::make_shared<IFFTOp>(ng_inputs.at(0), inverse, centered)};
+        return {std::make_shared<IFFTOp>(ng_inputs.at(0), ng_inputs.at(1), inverse, centered)};
     });
     ngraph::onnx_import::register_operator(ComplexMulOp::type_info.name, 1, "", [](const ngraph::onnx_import::Node& node) -> ngraph::OutputVector {
         ngraph::OutputVector ng_inputs {node.get_ng_inputs()};