def test_upsample_with_second_input_infer(self, scales, input_shape,
expected_shape):
nodes_attributes['scales'] = {'kind': 'data', 'value': scales}
graph = build_graph(
nodes_attributes, [('node_1', 'upsample'), ('scales', 'upsample'),
('upsample', 'node_3'),
('node_3', 'op_output')], {
'node_3': {
'shape': None
},
'node_1': {
'shape': input_shape
},
'upsample': {
'mode': 'linear',
'height_scale': None,
'width_scale': None
}
})
graph.graph['layout'] = 'NCHW'
upsample_node = Node(graph, 'upsample')
UpsampleOp.upsample_infer(upsample_node)
res_shape = graph.node['node_3']['shape']
for i in range(0, len(expected_shape)):
self.assertEqual(expected_shape[i], res_shape[i])
def test_upsample_with_scales_infer(self):
graph = build_graph(
nodes_attributes, [('node_1', 'upsample'), ('upsample', 'node_3'),
('node_3', 'op_output')], {
'node_3': {
'shape': None
},
'node_1': {
'shape': np.array([1, 3, 227, 227])
},
'upsample': {
'mode': 'linear',
'height_scale': 2.,
'width_scale': 2.
}
})
graph.graph['layout'] = 'NCHW'
upsample_node = Node(graph, 'upsample')
UpsampleOp.upsample_infer(upsample_node)
exp_shape = np.array([1, 3, 454, 454])
res_shape = graph.node['node_3']['shape']
for i in range(0, len(exp_shape)):
self.assertEqual(exp_shape[i], res_shape[i])
def test_upsample_with_scales_infer(self, scales, input_shape,
expected_shape):
graph = build_graph(
nodes_attributes, [('node_1', 'upsample'), ('upsample', 'node_3'),
('node_3', 'op_output')], {
'node_3': {
'shape': None,
'value': None
},
'node_1': {
'shape': input_shape
},
'upsample': {
'mode': 'linear',
'height_scale': scales[2],
'width_scale': scales[3]
}
})
graph.graph['layout'] = 'NCHW'
upsample_node = Node(graph, 'upsample')
UpsampleOp.upsample_infer(upsample_node)
res_shape = graph.node['node_3']['shape']
self.assertTrue(strict_compare_tensors(expected_shape, res_shape))