Exemple #1
0
    def test_accum_infer_have_reference(self):
        graph = build_graph(
            nodes_attributes, [('node_1', 'accum'), ('node_2', 'accum'),
                               ('accum', 'node_3'), ('node_3', 'op_output')], {
                                   'node_3': {
                                       'shape': None
                                   },
                                   'node_1': {
                                       'shape': np.array([1, 3, 227, 227])
                                   },
                                   'node_2': {
                                       'shape': np.array([1, 3, 227, 227])
                                   },
                                   'accum': {
                                       'top_height': 0,
                                       'top_width': 0,
                                       'size_divisible_by': 0,
                                       'have_reference': 1
                                   }
                               })

        accum_node = Node(graph, 'accum')
        AccumOp.accum_infer(accum_node)
        exp_shape = np.array([1, 6, 227, 227])
        res_shape = graph.node['node_3']['shape']
        for i in range(0, len(exp_shape)):
            self.assertEqual(exp_shape[i], res_shape[i])
Exemple #2
0
    def extract(cls, node):
        proto_layer = node.pb
        param = proto_layer.accum_param

        attrs = collect_attributes(param)
        # update the attributes of the node
        AccumOp.update_node_stat(node, attrs)
        return cls.enabled