def extract(cls, node):
pb = node.parameters
collect_until_token(pb, b'<PoolSize>')
kernel = read_binary_integer32_token(pb)
tag = find_next_tag(pb)
if tag == '<PoolStep>':
read_placeholder(pb, 1)
stride = read_binary_integer32_token(pb)
pool_step = stride
pool_stride = read_token_value(pb, b'<PoolStride>')
elif tag == '<PoolStride>':
stride = 1
pool_step = None
read_placeholder(pb, 1)
pool_stride = read_binary_integer32_token(pb)
else:
raise Error('Can not extract parameters for {}'.format(node))
mapping_rule = {
'window': np.array([1, 1, 1, kernel], dtype=np.int64),
'stride': np.array([1, 1, 1, stride], dtype=np.int64),
'pool_stride': pool_stride,
'pool_step': pool_step,
'pad': np.array([[0, 0], [0, 0], [0, 0], [0, 0]], dtype=np.int64),
'pad_spatial_shape': np.array([[0, 0], [0, 0]], dtype=np.int64),
'pool_method': 'max',
}
mapping_rule.update(layout_attrs())
Pooling.update_node_stat(node, mapping_rule)
return cls.enabled
def _insert_pooling(graph: Graph, first_node: Node, second_node: Node,
spatial_dims):
"""
This function inserts point wise pooling layer between two nodes
"""
log.debug("STRIDE PROP: Insert pooling between {} and {}".format(
first_node.name, second_node.name))
stride_prop = second_node.stride_prop
assert len(graph.get_edge_data(first_node.id, second_node.id)) == 1
eattrs = graph.get_edge_data(first_node.id, second_node.id)[0]
graph.remove_edge(first_node.id, second_node.id)
pooling = Pooling(
graph,
dict(name='Pooling_',
spatial_dims=spatial_dims,
window=np.array([1, 1, 1, 1]),
output_spatial_shape=None,
stride=np.array(stride_prop),
pad_spatial_shape=np.array([[0, 0], [0, 0]]),
pad=np.array([[0, 0], [0, 0], [0, 0], [0, 0]]),
pool_method='max',
is_partial_inferred=False))
pooling_data = pooling.create_node_with_data([first_node])
_clean_fw_tensor_attrs(pooling_data)
graph.add_edges_from([(pooling_data.id, second_node.id, eattrs)])
def extract(cls, node):
# Extract pads attribute
final_pads = get_pads(node.module)
# Extract strides attribute
strides = [node.module.stride, node.module.stride]
final_strides = np.array([1, 1, *strides], dtype=np.int64)
kernel_shape = [node.module.kernel_size, node.module.kernel_size]
final_kernel_shape = np.array([1, 1, *kernel_shape], dtype=np.int64)
attrs = {
'op': node.op,
'window': final_kernel_shape,
'stride': final_strides,
'pad': final_pads,
'pool_method': 'avg',
'exclude_pad': False,
'channel_dims': np.array([1], dtype=np.int64),
'batch_dims': np.array([0], dtype=np.int64),
'layout': 'NCHW',
}
# update the attributes of the node
Pooling.update_node_stat(node, attrs)
return cls.enabled
def infer(cls, node: Node):
input_shape = node.in_node(0).shape
input_h = input_shape[2]
input_w = input_shape[3]
output_h = node.output_size[0]
output_w = node.output_size[1]
stride_h = input_h // output_h
stride_w = input_w // output_w
kernel_h = input_h - (output_h - 1) * stride_h
kernel_w = input_w - (output_w - 1) * stride_w
data = {
'window': int64_array([1, 1, kernel_h, kernel_w]),
'stride': int64_array([1, 1, stride_h, stride_w]),
'pad': int64_array([[0, 0], [0, 0], [0, 0], [0, 0]]),
'pad_spatial_shape': int64_array([[0, 0], [0, 0]]),
'pool_method': 'avg',
'exclude_pad': False,
'output_spatial_shape': None,
'spatial_dims': None,
'channel_dims': int64_array([1]),
'batch_dims': int64_array([0]),
'layout': 'NCHW',
'rounding_type': 'floor',
'pooling_convention': 'valid'
}
# update the attributes of the node
Pooling.update_node_stat(node, data)
Pooling.infer(node)
def test_pooling_infer_wrong_input_shape(self):
graph = build_graph(
nodes_attributes, [('node_1', 'pool'), ('pool', 'node_2'),
('node_2', 'op_output')],
{
'node_2': {
'shape': None
},
'node_1': {
'shape': np.array([1, 3, 1, 1])
},
'pool': {
'window': np.array([1, 1, 5, 5]),
'stride': np.array([1, 1, 2, 2]),
'pad': np.array([[0, 0], [0, 0], [1, 1], [1, 1]]),
'pad_spatial_shape': np.array([[1, 1], [1, 1]]),
'pool_method': 'avg',
'exclude_pad': False,
'global_pool': False,
'output_spatial_shape': None,
'output_shape': None,
'kernel_spatial': np.array([3, 3]),
'spatial_dims': np.array([2, 3]),
'channel_dims': np.array([1]),
'batch_dims': np.array([0]),
'pooling_convention': 'full'
}
})
pool_node = Node(graph, 'pool')
with self.assertRaises(Error):
Pooling.infer(pool_node)
def test_pooling_infer_no_shape(self):
graph = build_graph(
nodes_attributes, [('node_1', 'pool'), ('pool', 'node_2'),
('node_2', 'op_output')],
{
'node_2': {
'shape': None
},
'node_1': {
'shape': None
},
'pool': {
'window': np.array([1, 1, 1, 1]),
'stride': np.array([1, 1, 2, 2]),
'pad': np.array([[0, 0], [0, 0], [3, 3], [3, 3]]),
'pad_spatial_shape': np.array([[3, 3], [3, 3]]),
'pool_method': 'avg',
'exclude_pad': False,
'output_spatial_shape': None,
'output_shape': None,
'kernel_spatial': np.array([3, 3]),
'spatial_dims': np.array([2, 3]),
'channel_dims': np.array([1]),
'batch_dims': np.array([0]),
'pooling_convention': 'full'
}
})
pool_node = Node(graph, 'pool')
Pooling.infer(pool_node)
res_shape = graph.node['node_2']['shape']
self.assertIsNone(res_shape)
def test_pooling_infer_no_convention(self):
graph = build_graph(
nodes_attributes, [('node_1', 'pool'), ('pool', 'node_2'),
('node_2', 'op_output')],
{
'node_2': {
'shape': None
},
'node_1': {
'shape': np.array([1, 3, 256, 256])
},
'pool': {
'window': np.array([1, 1, 1, 1]),
'stride': np.array([1, 1, 2, 2]),
'pad': np.array([[0, 0], [0, 0], [3, 3], [3, 3]]),
'pad_spatial_shape': np.array([[3, 3], [3, 3]]),
'pool_method': 'avg',
'exclude_pad': False,
'global_pool': False,
'output_spatial_shape': None,
'output_shape': None,
'kernel_spatial': np.array([3, 3]),
'spatial_dims': np.array([2, 3]),
'channel_dims': np.array([1]),
'batch_dims': np.array([0])
}
})
pool_node = Node(graph, 'pool')
Pooling.infer(pool_node)
exp_shape = np.array([1, 3, 130, 130])
res_shape = graph.node['node_2']['shape']
for i in range(0, len(exp_shape)):
self.assertEqual(exp_shape[i], res_shape[i])
def extract(cls, node):
attrs = common_onnx_pool_extractor(node)
attrs.update({
'pooling_convention': 'full',
'global_pool': True,
})
Pooling.update_node_stat(node, attrs)
return cls.enabled
def extract(cls, node):
proto_layer = node.pb
param = proto_layer.pooling_param
method = 'max'
exclude_pad = True
kernel = [0, 0]
stride = [1, 1]
padding = [0, 0]
global_pooling = False
if hasattr(param, 'global_pooling') and param.global_pooling:
global_pooling = param.global_pooling
else:
kernel = get_spatial_attr(kernel, 'kernel_size', 'kernel', param)
padding = get_spatial_attr(padding, 'pad', 'pad', param)
stride = get_spatial_attr(stride, 'stride', 'stride', param)
if param.pool == 0:
method = 'max'
exclude_pad = True
elif param.pool == 1:
method = 'avg'
exclude_pad = False
else:
raise ValueError('Unknown Pooling Method!')
pooling_convention = 'full' # for Caffe rounding type should be ceil
rt = 'ceil'
if hasattr(param, 'ceil_mode') and not param.ceil_mode:
# If pooling has ceil_mode and ceil_mode is False using floor for rounding shapes in partial_infer
pooling_convention = 'valid'
rt = 'floor'
attrs = {
'window': int64_array([1, 1, kernel[1], kernel[0]]),
'stride': int64_array([1, 1, stride[1], stride[0]]),
'pad': int64_array([[0, 0], [0, 0], [padding[1], padding[1]], [padding[0], padding[0]]]),
'pad_spatial_shape': int64_array([[padding[1], padding[1]], [padding[0], padding[0]]]),
'pool_method': method,
'exclude_pad': exclude_pad,
'global_pool': global_pooling,
'output_spatial_shape': None,
'rounding_type': rt
}
attrs.update(layout_attrs())
attrs['pooling_convention'] = pooling_convention
# update the attributes of the node
Pooling.update_node_stat(node, attrs)
return cls.enabled
def extract(cls, node):
attrs = get_mxnet_layer_attrs(node.symbol_dict)
kernel = attrs.tuple("kernel", int, None)
stride = attrs.tuple("stride", int,
tuple(np.ones(len(kernel), dtype=np.int64)))
padding = attrs.tuple("pad", int,
tuple(np.zeros(len(kernel), dtype=np.int64)))
method = attrs.str("pool_type", None)
rt = 'floor'
data = {
'window':
np.array([1, 1, *[k for k in kernel]], dtype=np.int64),
'stride':
np.array([1, 1, *[s for s in stride]], dtype=np.int64),
'pad':
np.array([[0, 0], [0, 0], *[[pad, pad] for pad in padding]],
dtype=np.int64),
'pad_spatial_shape':
np.array([[pad, pad] for pad in padding], dtype=np.int64),
'pool_method':
method,
'exclude_pad':
False,
'output_spatial_shape':
None,
'spatial_dims':
None,
'channel_dims':
np.array([1], dtype=np.int64),
'batch_dims':
np.array([0], dtype=np.int64),
'layout':
'NCHW',
'rounding_type':
rt,
}
pooling_conv = attrs.str("pooling_convention", 'valid')
if pooling_conv:
data["pooling_convention"] = pooling_conv
if pooling_conv == 'full':
data["rounding_type"] = 'ceil'
global_pool = attrs.bool("global_pool", False)
if global_pool:
data["global_pool"] = global_pool
# update the attributes of the node
Pooling.update_node_stat(node, data)
return cls.enabled
def test_pooling_dynamic_infer(self):
graph = build_graph(
nodes_attributes, [('node_1', 'pool'), ('pool', 'node_2'),
('node_2', 'op_output')],
{
'node_2': {
'shape': None
},
'node_1': {
'shape':
shape_array([
1, dynamic_dimension_value, dynamic_dimension_value,
256
])
},
'pool': {
'window': np.array([1, 1, 1, 1]),
'stride': np.array([1, 1, 2, 2]),
'pad': np.array([[0, 0], [0, 0], [3, 3], [3, 3]]),
'pad_spatial_shape': np.array([[3, 3], [3, 3]]),
'pool_method': 'avg',
'exclude_pad': False,
'global_pool': False,
'output_spatial_shape': None,
'output_shape': None,
'kernel_spatial': np.array([3, 3]),
'spatial_dims': np.array([2, 3]),
'channel_dims': np.array([1]),
'batch_dims': np.array([0]),
'pooling_convention': 'full'
}
})
pool_node = Node(graph, 'pool')
Pooling.infer(pool_node)
exp_shape = shape_array(
[1, dynamic_dimension_value, dynamic_dimension_value, 131])
res_shape = graph.node['node_2']['shape']
self.assertTrue(strict_compare_tensors(exp_shape, res_shape))
def find_and_replace_pattern(self, graph: Graph):
for pool_v2_node in graph.get_op_nodes(op='PoolingV2'):
pool_v2_name = pool_v2_node.soft_get('name', pool_v2_node.id)
pool_v1_node = Pooling(
graph, {
'window': pool_v2_node.in_port(1).data.get_value(),
'stride': pool_v2_node.in_port(2).data.get_value(),
'pad': pool_v2_node.pad,
'spatial_dims': pool_v2_node.spatial_dims,
'auto_pad': pool_v2_node.auto_pad,
'output_spatial_shape': pool_v2_node.output_spatial_shape,
'pad_spatial_shape': pool_v2_node.pad_spatial_shape,
'pool_method': pool_v2_node.pool_method,
'permute_attrs': pool_v2_node.permute_attrs,
}).create_node()
rename_nodes([(pool_v2_node, pool_v2_name + '/to_be_removed'),
(pool_v1_node, pool_v2_name)])
pool_v2_node.in_port(0).get_connection().set_destination(
pool_v1_node.in_port(0))
pool_v2_node.out_port(0).get_connection().set_source(
pool_v1_node.out_port(0))
def extract(cls, node):
attrs = common_onnx_pool_extractor(node)
Pooling.update_node_stat(node, attrs)
return cls.enabled
def extract(cls, node):
attrs = create_pooling_attrs(node, 'avg')
attrs.update({'op': __class__.op})
# update the attributes of the node
Pooling.update_node_stat(node, attrs)
return cls.enabled
