Skip to content

Commit 7c7cc34

Browse files
nanoskriptnanoskript
committed
feat: add support for ExtractImagePatches
Signed-off-by: Nanoskript <96655713+nanoskript@users.noreply.github.com>
1 parent 0004163 commit 7c7cc34

File tree

5 files changed

+187
-0
lines changed

5 files changed

+187
-0
lines changed

‎tests/test_backend.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
matrix_diag_part = tf.compat.v1.matrix_diag_part
7474
fake_quant_with_min_max_args = tf.quantization.fake_quant_with_min_max_args
7575
fake_quant_with_min_max_vars = tf.quantization.fake_quant_with_min_max_vars
76+
extract_image_patches = tf.image.extract_patches
7677
elif Version(tf.__version__) >= Version("1.13"):
7778
conv2d_backprop_input = tf.compat.v1.nn.conv2d_backprop_input
7879
conv3d_transpose = tf.compat.v1.nn.conv3d_transpose
@@ -96,6 +97,7 @@
9697
matrix_diag_part = tf.compat.v1.matrix_diag_part
9798
fake_quant_with_min_max_args = tf.compat.v1.quantization.fake_quant_with_min_max_args
9899
fake_quant_with_min_max_vars = tf.compat.v1.quantization.fake_quant_with_min_max_vars
100+
extract_image_patches = tf.compat.v1.extract_image_patches
99101
else:
100102
conv2d_backprop_input = tf.nn.conv2d_backprop_input
101103
conv3d_transpose = tf.nn.conv3d_transpose
@@ -113,6 +115,7 @@
113115
is_inf = tf.is_inf
114116
floormod = tf.floormod
115117
matrix_diag_part = tf.matrix_diag_part
118+
extract_image_patches = tf.extract_image_patches
116119

117120

118121
def make_xval(shape):
@@ -6361,5 +6364,22 @@ def func(tensor, indices, updates):
63616364
self._run_test_case(func, [_OUTPUT], {_INPUT: tensor_val, _INPUT1: indices_val, _INPUT2: updates_val})
63626365
self._run_test_case(func, [_OUTPUT], {_INPUT: tensor_val, _INPUT1: indices64_val, _INPUT2: updates_val})
63636366

6367+
@check_opset_min_version(9, "EyeLike and ConstantOfShape")
6368+
def test_extract_image_patches(self):
6369+
for rates in [[1, 1], [1, 4], [4, 1], [3, 3]]:
6370+
for _, padding, x_shape, sizes, strides in get_conv_getdata():
6371+
def func(x):
6372+
return extract_image_patches(
6373+
x,
6374+
sizes=sizes,
6375+
strides=strides,
6376+
rates=[1] + rates + [1],
6377+
padding=padding,
6378+
name=_TFOUTPUT
6379+
)
6380+
6381+
x_val = make_xval(x_shape)
6382+
self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
6383+
63646384
if __name__ == '__main__':
63656385
unittest_main()

‎tf2onnx/onnx_opset/nn.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2091,3 +2091,81 @@ def version_11(cls, ctx, node, **kwargs):
20912091
ctx.replace_all_inputs(node.output[3], sum_max_neg)
20922092

20932093
ctx.remove_node(node.name)
2094+
2095+
2096+
@tf_op("ExtractImagePatches")
2097+
class ExtractImagePatches:
2098+
@classmethod
2099+
def version_9(cls, ctx, node, **kwargs):
2100+
input_shape = ctx.get_shape(node.input[0])
2101+
output_shape = node.output_shapes[0]
2102+
2103+
sizes = node.get_attr_value("ksizes")
2104+
strides = node.get_attr_value("strides")
2105+
rates = node.get_attr_value("rates")
2106+
padding = node.get_attr_str("padding")
2107+
2108+
# Our constraints.
2109+
utils.make_sure(0 not in output_shape, "Empty ExtractImagePatches output is unsupported.")
2110+
[_, size_rows, size_cols, _] = sizes
2111+
2112+
# Transform input into [N * C, H, W, 1].
2113+
transformed_input = ctx.make_node("Reshape", inputs=[
2114+
ctx.make_node("Transpose", inputs=node.input, attr=dict(perm=[0, 3, 1, 2])).output[0],
2115+
ctx.make_const(utils.make_name("new_shape"), np.int64([
2116+
input_shape[0] * input_shape[3],
2117+
input_shape[1],
2118+
input_shape[2],
2119+
1,
2120+
])).output[0],
2121+
])
2122+
2123+
# Create identity kernel.
2124+
k = size_rows * size_cols
2125+
identity_kernel = ctx.make_node("Reshape", inputs=[
2126+
ctx.make_node("EyeLike", inputs=[
2127+
ctx.make_node("ConstantOfShape", inputs=[
2128+
ctx.make_const(utils.make_name("eye_size"), np.array([k, k], dtype=np.int64)).output[0],
2129+
]).output[0],
2130+
]).output[0],
2131+
ctx.make_const(utils.make_name("new_shape"), np.array([
2132+
size_rows,
2133+
size_cols,
2134+
1,
2135+
k,
2136+
], dtype=np.int64)).output[0],
2137+
])
2138+
2139+
# Construct placeholder convolution node and transform into [N * C, K, ?H, ?W].
2140+
convolution = ctx.make_node("Conv", inputs=[transformed_input.output[0], identity_kernel.output[0]],
2141+
shapes=[[input_shape[0] * input_shape[3], output_shape[1], output_shape[2], k]],
2142+
attr=dict(strides=strides, dilations=rates, padding=padding, data_format="NHWC"),
2143+
dtypes=node.output_dtypes)
2144+
2145+
# Transform into [N, ?H, ?W, C * K].
2146+
output_node = ctx.make_node("Reshape", inputs=[
2147+
ctx.make_node("Transpose", inputs=[
2148+
ctx.make_node("Reshape", inputs=[
2149+
convolution.output[0],
2150+
ctx.make_const(utils.make_name("new_shape"), np.array([
2151+
input_shape[0],
2152+
input_shape[3],
2153+
output_shape[1],
2154+
output_shape[2],
2155+
k,
2156+
], dtype=np.int64)).output[0],
2157+
]).output[0],
2158+
], attr=dict(perm=[0, 2, 3, 4, 1])).output[0],
2159+
ctx.make_const(utils.make_name("new_shape"), np.array(output_shape, dtype=np.int64)).output[0],
2160+
])
2161+
2162+
# Replace original node.
2163+
ctx.replace_all_inputs(node.output[0], output_node.output[0])
2164+
ctx.remove_node(node.name)
2165+
2166+
# Transform convolution node.
2167+
kernel_shape = conv_kernel_shape(ctx, convolution, 1)
2168+
strides = conv_dims_attr(convolution, "strides")
2169+
dilations = conv_dims_attr(convolution, "dilations")
2170+
add_padding(ctx, convolution, kernel_shape, strides, dilations)
2171+
conv_convert_inputs(ctx, convolution, with_kernel=True)

‎tf2onnx/rewriter/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from tf2onnx.rewriter.lstm_tf2_rewriter import rewriter_lstm_tf2
2525
from tf2onnx.rewriter.gru_tf2_rewriter import rewrite_gru_tf2
2626
from tf2onnx.rewriter.fused_op_rewriter import rewrite_fused_ops
27+
from tf2onnx.rewriter.extract_image_patches_rewriter import rewrite_extract_image_patches
2728

2829

2930
__all__ = [
@@ -53,4 +54,5 @@
5354
"rewriter_lstm_tf2",
5455
"rewrite_gru_tf2",
5556
"rewrite_fused_ops",
57+
"rewrite_extract_image_patches",
5658
]

‎tf2onnx/rewriter/extract_image_patches_rewriter.py

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
3+
4+
"""
5+
tf2onnx.rewriter.extract_image_patches_rewriter - Rewrites ExtractImagePatches into supported operations.
6+
"""
7+
8+
import numpy as np
9+
from tf2onnx import utils
10+
from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
11+
12+
13+
# pylint: disable=missing-docstring
14+
15+
def rewrite_extract_image_patches(g, ops):
16+
pattern = OpTypePattern("ExtractImagePatches", name="extract_image_patches")
17+
matcher = GraphMatcher(pattern)
18+
match_results = list(matcher.match_ops(ops))
19+
for match_result in match_results:
20+
operation = match_result.get_op("extract_image_patches")
21+
input_shape = g.get_shape(operation.input[0])
22+
output_shape = operation.output_shapes[0]
23+
24+
sizes = operation.get_attr_value("ksizes")
25+
strides = operation.get_attr_value("strides")
26+
rates = operation.get_attr_value("rates")
27+
padding = operation.get_attr_str("padding")
28+
29+
# Our constraints.
30+
utils.make_sure(0 not in output_shape, "Empty ExtractImagePatches output is unsupported.")
31+
[_, size_rows, size_cols, _] = sizes
32+
33+
# Transform input into [N * C, H, W, 1].
34+
transformed_input = g.make_node("Reshape", inputs=[
35+
g.make_node("Transpose", inputs=operation.input, attr=dict(perm=[0, 3, 1, 2])).output[0],
36+
g.make_const(utils.make_name("new_shape"), np.int64([
37+
input_shape[0] * input_shape[3],
38+
input_shape[1],
39+
input_shape[2],
40+
1,
41+
])).output[0],
42+
])
43+
44+
# Create identity kernel.
45+
k = size_rows * size_cols
46+
identity_kernel = g.make_node("Reshape", inputs=[
47+
g.make_node("EyeLike", inputs=[
48+
g.make_node("ConstantOfShape", inputs=[
49+
g.make_const(utils.make_name("eye_size"), np.array([k, k], dtype=np.int64)).output[0],
50+
]).output[0],
51+
]).output[0],
52+
g.make_const(utils.make_name("new_shape"), np.array([
53+
size_rows,
54+
size_cols,
55+
1,
56+
k,
57+
], dtype=np.int64)).output[0],
58+
])
59+
60+
# Convolve into [N * C, ?H, ?W, K].
61+
convolution = g.make_node("Conv2D", inputs=[transformed_input.output[0], identity_kernel.output[0]],
62+
attr=dict(strides=strides, dilations=rates, padding=padding, data_format="NHWC"),
63+
shapes=[[input_shape[0] * input_shape[3], output_shape[1], output_shape[2], k]],
64+
dtypes=operation.output_dtypes, skip_conversion=False)
65+
66+
# Transform into [N, ?H, ?W, C * K].
67+
output_node = g.make_node("Reshape", inputs=[
68+
g.make_node("Transpose", inputs=[
69+
g.make_node("Reshape", inputs=[
70+
convolution.output[0],
71+
g.make_const(utils.make_name("new_shape"), np.array([
72+
input_shape[0],
73+
input_shape[3],
74+
output_shape[1],
75+
output_shape[2],
76+
k,
77+
], dtype=np.int64)).output[0],
78+
]).output[0],
79+
], attr=dict(perm=[0, 2, 3, 4, 1])).output[0],
80+
g.make_const(utils.make_name("new_shape"), np.array(output_shape, dtype=np.int64)).output[0],
81+
])
82+
83+
# Replace node.
84+
g.replace_all_inputs(operation.output[0], output_node.output[0])
85+
g.remove_node(operation.name)
86+
return g.get_nodes()

‎tf2onnx/tfonnx.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -598,6 +598,7 @@ def compat_handler(ctx, node, **kwargs):
598598
rewriter_lstm_tf2,
599599
rewrite_gru_tf2,
600600
rewrite_single_direction_lstm,
601+
rewrite_extract_image_patches,
601602
# bi-directional
602603
rewrite_bi_direction_lstm,
603604
rewrite_single_direction_gru,

0 commit comments

Comments
 (0)