def extract(cls, node):
attrs = {
'axis': 1,
'end_axis': -1,
}
Flatten.update_node_stat(node, attrs)
return cls.enabled
def extract(node):
attrs = {
'axis': 1,
'end_axis': -1,
}
Flatten.update_node_stat(node, attrs)
return __class__.enabled
def extract(cls, node):
proto_layer = node.pb
param = proto_layer.flatten_param
attrs = {
'axis': param.axis,
'end_axis': param.end_axis,
}
Flatten.update_node_stat(node, attrs)
return cls.enabled
def test_flatten_infer_no_shape(self):
graph = build_graph(
nodes_attributes, [('node_1', 'flatten_1'),
('flatten_1', 'node_2')], {
'node_2': {
'is_output': True,
'shape': None
},
'node_1': {
'shape': None
},
'flatten_1': {
'axis': 1
}
})
flatten_node = Node(graph, 'flatten_1')
Flatten.infer(flatten_node)
res_shape = graph.node['node_2']['shape']
self.assertIsNone(res_shape)
def test_flatten_infer(self):
graph = build_graph(
nodes_attributes, [('node_1', 'flatten_1'),
('flatten_1', 'node_2')], {
'node_2': {
'is_output': True,
'shape': np.array([1, 3 * 256 * 256])
},
'node_1': {
'shape': np.array([1, 3, 256, 256])
},
'flatten_1': {
'axis': 1,
'dim': []
}
})
flatten_node = Node(graph, 'flatten_1')
Flatten.infer(flatten_node)
exp_shape = np.array([1, 3 * 256 * 256])
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 = {
'axis': node.module.axis,
}
Flatten.update_node_stat(node, attrs)
return cls.enabled