Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings
This repository was archived by the owner on May 29, 2023. It is now read-only.

Commit a82c0d7

Browse files
authored
Add FFT (#8)
Support torch.fft
1 parent f4c126e commit a82c0d7

File tree

18 files changed

+405
-86
lines changed

18 files changed

+405
-86
lines changed

‎.github/workflows/main.yml‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -44,19 +44,32 @@ jobs:
4444
mkdir build && cd build
4545
cmake .. -DCMAKE_BUILD_TYPE=Release && make -j$(nproc --all)
4646
47-
- name: Test default parameters
47+
- name: (Test Unpool) default parameters
4848
run: |
4949
source /opt/intel/openvino/bin/setupvars.sh
50-
python3 export_model.py --mode=default
51-
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py --input_model model_with_unpool.onnx --extension mo_extensions
50+
python3 examples/unpool/export_model.py --mode=default
51+
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py --input_model model.onnx --extension mo_extensions
5252
python3 compare.py
5353
54-
- name: Test dynamic shapes
54+
- name: (Test Unpool) Test dynamic shapes
5555
run: |
5656
source /opt/intel/openvino/bin/setupvars.sh
57-
python3 export_model.py --mode=dynamic_size --shape 5 3 6 9
58-
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py --input_model model_with_unpool.onnx --extension mo_extensions
57+
python3 examples/unpool/export_model.py --mode=dynamic_size --shape 5 3 6 9
58+
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py --input_model model.onnx --extension mo_extensions
5959
python3 compare.py
6060
61-
python3 export_model.py --mode=dynamic_size --shape 4 3 17 8
61+
python3 examples/unpool/export_model.py --mode=dynamic_size --shape 4 3 17 8
62+
python3 compare.py
63+
64+
65+
- name: (Test FFT)
66+
run: |
67+
source /opt/intel/openvino/bin/setupvars.sh
68+
for shape in "5 120 2" "4 240 320 2" "3 5 240 320 2"; do
69+
python3 examples/fft/export_model.py --shape ${shape}
70+
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py --input_model model.onnx --extension mo_extensions
71+
python3 compare.py
72+
done
73+
74+
python3 examples/fft/export_model.py --shape 3 7 140 100 2
6275
python3 compare.py

‎README.md‎

Lines changed: 9 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,24 @@
1-
Guide of how to enable PyTorch `nn.MaxUnpool2d`in Intel OpenVINO.
1+
Repository with guides to enable some layers from PyTorch in Intel OpenVINO:
22

3-
[![CI](https://github.com/dkurt/openvino_pytorch_unpool/workflows/CI/badge.svg?branch=master)](https://github.com/dkurt/openvino_pytorch_unpool/actions?query=branch%3Amaster)
3+
[![CI](https://github.com/dkurt/openvino_pytorch_layers/workflows/CI/badge.svg?branch=master)](https://github.com/dkurt/openvino_pytorch_layers/actions?query=branch%3Amaster)
44

5-
## Description
6-
There are two problems with OpenVINO and MaxUnpool at the moment of this guide creation:
5+
*[nn.MaxUnpool2d](examples/unpool)
6+
*[torch.fft](examples/fft)
77

8-
* OpenVINO does not have Unpooling kernels
9-
* PyTorch -> ONNX conversion is unimplemented for `nn.MaxUnpool2d`
10-
11-
So following this guide you will learn
12-
* How to perform PyTorch -> ONNX conversion for unsupported layers
13-
* How to convert ONNX to OpenVINO Intermediate Respresentation (IR) with extensions
14-
* How to write custom CPU layers in OpenVINO
15-
16-
## Get ONNX model
17-
18-
MaxUnpool layer in PyTorch takes two inputs - input `features` from any layer and `indices` after MaxPool layer:
19-
20-
```python
21-
self.pool = nn.MaxPool2d(2, stride=2, return_indices=True)
22-
self.unpool = nn.MaxUnpool2d(2, stride=2)
23-
24-
output, indices = self.pool(x)
25-
# ...
26-
unpooled = self.unpool(features, indices)
27-
```
28-
29-
If your version of PyTorch does not support ONNX model conversion with MaxUnpool, replace every unpool layer definition
30-
```python
31-
self.unpool = nn.MaxUnpool2d(2, stride=2)
32-
```
33-
to
34-
```python
35-
self.unpool = Unpool2d()
36-
```
37-
38-
where `Unpool2d` defined in [unpool.py](./unpool.py). Also, replace op usage from
39-
40-
```python
41-
self.unpool(features, indices)
42-
```
43-
to
44-
```python
45-
self.unpool.apply(features, indices)
46-
```
47-
48-
See complete example in [export_model.py](./export_model.py).
498

509
## OpenVINO Model Optimizer extension
5110

52-
To create OpenVINO IR, use extra `--extension` flag to specify a path to Model Optimizer extensions that perform graph transformations and register unpooling layer.
11+
To create OpenVINO IR, use extra `--extension` flag to specify a path to Model Optimizer extensions that perform graph transformations and register custom layers.
5312

5413
```bash
5514
python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py \
56-
--input_model model_with_unpool.onnx \
57-
--extension mo_extensions
15+
--input_model model.onnx \
16+
--extension openvino_pytorch_layers/mo_extensions
5817
```
5918

6019
## Custom CPU extensions
6120

62-
You also need to build CPU extensions library which actually has Unpooling implementation:
21+
You also need to build CPU extensions library which actually has C++ layers implementations:
6322
```bash
6423
source /opt/intel/openvino/bin/setupvars.sh
6524
export TBB_DIR=/opt/intel/openvino/deployment_tools/inference_engine/external/tbb/cmake/
@@ -77,6 +36,6 @@ from openvino.inference_engine import IECore
7736
ie = IECore()
7837
ie.add_extension('user_ie_extensions/build/libuser_cpu_extension.so', 'CPU')
7938

80-
net = ie.read_network('model_with_unpool.xml', 'model_with_unpool.bin')
39+
net = ie.read_network('model.xml', 'model.bin')
8140
exec_net = ie.load_network(net, 'CPU')
8241
```

‎compare.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
ie = IECore()
88
ie.add_extension('user_ie_extensions/build/libuser_cpu_extension.so', 'CPU')
99

10-
net = ie.read_network('model_with_unpool.xml', 'model_with_unpool.bin')
10+
net = ie.read_network('model.xml', 'model.bin')
1111
net.reshape({'input': inp.shape})
1212
exec_net = ie.load_network(net, 'CPU')
1313

@@ -17,5 +17,5 @@
1717
maxdiff = np.max(np.abs(ref - out))
1818
print('Reference range: [{}, {}]'.format(np.min(ref), np.max(ref)))
1919
print('Maximal difference:', maxdiff)
20-
if maxdiff > 1e-6:
20+
if maxdiff > 1e-5:
2121
exit(1)

‎examples/fft/export_model.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import numpy as np
2+
import argparse
3+
import torch
4+
import torch.nn as nn
5+
from torch.autograd import Variable
6+
from fft import FFT
7+
8+
np.random.seed(324)
9+
torch.manual_seed(32)
10+
11+
class MyModel(nn.Module):
12+
def __init__(self):
13+
super(MyModel, self).__init__()
14+
self.fft = FFT()
15+
16+
def forward(self, x):
17+
y = self.fft.apply(x, False)
18+
y = y * 2
19+
y = self.fft.apply(y, True)
20+
return y
21+
22+
parser = argparse.ArgumentParser(description='Generate ONNX model and test data')
23+
parser.add_argument('--shape', type=int, nargs='+', default=[5, 3, 6, 8, 2])
24+
args = parser.parse_args()
25+
26+
model = MyModel()
27+
inp = Variable(torch.randn(args.shape))
28+
model.eval()
29+
30+
with torch.no_grad():
31+
torch.onnx.export(model, inp, 'model.onnx',
32+
input_names=['input'],
33+
output_names=['output'],
34+
operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
35+
36+
ref = model(inp)
37+
np.save('inp', inp.detach().numpy())
38+
np.save('ref', ref.detach().numpy())

‎examples/fft/fft.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import torch
2+
3+
class FFT(torch.autograd.Function):
4+
@staticmethod
5+
def symbolic(g, x, inverse):
6+
return g.op('FFT', x,
7+
inverse_i=inverse)
8+
9+
@staticmethod
10+
def forward(self, x, inverse):
11+
# https://pytorch.org/docs/stable/torch.html#torch.fft
12+
signal_ndim = 2 if len(x.shape) == 5 else 1
13+
if inverse:
14+
y = torch.ifft(input=x, signal_ndim=signal_ndim, normalized=True)
15+
else:
16+
y = torch.fft(input=x, signal_ndim=signal_ndim, normalized=True)
17+
return y

‎examples/unpool/README.md‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
Guide of how to enable PyTorch `nn.MaxUnpool2d` in Intel OpenVINO.
2+
3+
4+
## Description
5+
There are two problems with OpenVINO and MaxUnpool at the moment of this guide creation:
6+
7+
* OpenVINO does not have Unpooling kernels
8+
* PyTorch -> ONNX conversion is unimplemented for `nn.MaxUnpool2d`
9+
10+
So following this guide you will learn
11+
* How to perform PyTorch -> ONNX conversion for unsupported layers
12+
* How to convert ONNX to OpenVINO Intermediate Respresentation (IR) with extensions
13+
* How to write custom CPU layers in OpenVINO
14+
15+
## Get ONNX model
16+
17+
MaxUnpool layer in PyTorch takes two inputs - input `features` from any layer and `indices` after MaxPool layer:
18+
19+
```python
20+
self.pool = nn.MaxPool2d(2, stride=2, return_indices=True)
21+
self.unpool = nn.MaxUnpool2d(2, stride=2)
22+
23+
output, indices = self.pool(x)
24+
# ...
25+
unpooled = self.unpool(features, indices)
26+
```
27+
28+
If your version of PyTorch does not support ONNX model conversion with MaxUnpool, replace every unpool layer definition
29+
```python
30+
self.unpool = nn.MaxUnpool2d(2, stride=2)
31+
```
32+
to
33+
```python
34+
self.unpool = Unpool2d()
35+
```
36+
37+
where `Unpool2d` defined in [unpool.py](./unpool.py). Also, replace op usage from
38+
39+
```python
40+
self.unpool(features, indices)
41+
```
42+
to
43+
```python
44+
self.unpool.apply(features, indices)
45+
```
46+
47+
See complete example in [export_model.py](./export_model.py).

‎export_model.py‎ renamed to ‎examples/unpool/export_model.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def forward(self, x):
3939
model.eval()
4040

4141
with torch.no_grad():
42-
torch.onnx.export(model, inp, 'model_with_unpool.onnx',
42+
torch.onnx.export(model, inp, 'model.onnx',
4343
input_names=['input'],
4444
output_names=['output'],
4545
operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
File renamed without changes.

‎mo_extensions/front/onnx/fft_ext.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
from mo.front.extractor import FrontExtractorOp
2+
from mo.front.onnx.extractors.utils import onnx_attr
3+
from ...ops.FFT import FFT
4+
5+
class FFTFrontExtractor(FrontExtractorOp):
6+
op = 'FFT'
7+
enabled = True
8+
9+
@classmethod
10+
def extract(cls, node):
11+
data = {
12+
'inverse': onnx_attr(node, 'inverse', 'i'),
13+
}
14+
15+
FFT.update_node_stat(node, data)
16+
return cls.enabled

‎mo_extensions/ops/FFT.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
from mo.front.common.partial_infer.elemental import copy_shape_infer
2+
from mo.graph.graph import Graph
3+
from mo.ops.op import Op
4+
5+
6+
class FFT(Op):
7+
op = 'FFT'
8+
enabled = True
9+
10+
def __init__(self, graph: Graph, attrs: dict):
11+
super().__init__(graph, {
12+
'type': __class__.op,
13+
'op': __class__.op,
14+
'in_ports_count': 1,
15+
'out_ports_count': 1,
16+
'infer': copy_shape_infer
17+
}, attrs)
18+
19+
def supported_attrs(self):
20+
return ['inverse']

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /