Python onnx.helper.make_attribute() Examples

def test_attr_int(self):  # type: () -> None
        # integer
        attr = helper.make_attribute("int", 3)
        self.assertEqual(attr.name, "int")
        self.assertEqual(attr.i, 3)
        checker.check_attribute(attr)
        # long integer
        attr = helper.make_attribute("int", 5)
        self.assertEqual(attr.name, "int")
        self.assertEqual(attr.i, 5)
        checker.check_attribute(attr)
        # octinteger
        attr = helper.make_attribute("int", 0o1701)
        self.assertEqual(attr.name, "int")
        self.assertEqual(attr.i, 0o1701)
        checker.check_attribute(attr)
        # hexinteger
        attr = helper.make_attribute("int", 0x1701)
        self.assertEqual(attr.name, "int")
        self.assertEqual(attr.i, 0x1701)
        checker.check_attribute(attr) 

def test_attr_string(self):  # type: () -> None
        # bytes
        attr = helper.make_attribute("str", b"test")
        self.assertEqual(attr.name, "str")
        self.assertEqual(attr.s, b"test")
        checker.check_attribute(attr)
        # unspecified
        attr = helper.make_attribute("str", "test")
        self.assertEqual(attr.name, "str")
        self.assertEqual(attr.s, b"test")
        checker.check_attribute(attr)
        # unicode
        attr = helper.make_attribute("str", u"test")
        self.assertEqual(attr.name, "str")
        self.assertEqual(attr.s, b"test")
        checker.check_attribute(attr) 

def test_attr_repeated_tensor_proto(self):  # type: () -> None
        tensors = [
            helper.make_tensor(
                name='a',
                data_type=TensorProto.FLOAT,
                dims=(1,),
                vals=np.ones(1).tolist()
            ),
            helper.make_tensor(
                name='b',
                data_type=TensorProto.FLOAT,
                dims=(1,),
                vals=np.ones(1).tolist()
            )]
        attr = helper.make_attribute("tensors", tensors)
        self.assertEqual(attr.name, "tensors")
        self.assertEqual(list(attr.tensors), tensors)
        checker.check_attribute(attr) 

def set_nodeattr(self, name, value):
        """Set a node attribute by name. Data is stored inside the ONNX node's
        AttributeProto container. Attribute must be part of get_nodeattr_types."""
        try:
            (dtype, req, def_val) = self.get_nodeattr_types()[name]
            attr = get_by_name(self.onnx_node.attribute, name)
            if attr is not None:
                # dtype indicates which ONNX Attribute member to use
                # (such as i, f, s...)
                if dtype == "s":
                    # encode string attributes
                    value = value.encode("utf-8")
                attr.__setattr__(dtype, value)
            else:
                # not set, create and insert AttributeProto
                attr_proto = helper.make_attribute(name, value)
                self.onnx_node.attribute.append(attr_proto)
        except KeyError:
            raise AttributeError("Op has no such attribute: " + name) 

def test_attr_float(self):  # type: () -> None
        # float
        attr = helper.make_attribute("float", 1.)
        self.assertEqual(attr.name, "float")
        self.assertEqual(attr.f, 1.)
        checker.check_attribute(attr)
        # float with scientific
        attr = helper.make_attribute("float", 1e10)
        self.assertEqual(attr.name, "float")
        self.assertEqual(attr.f, 1e10)
        checker.check_attribute(attr) 

def test_attr_doc_string(self):  # type: () -> None
        attr = helper.make_attribute("a", "value")
        self.assertEqual(attr.name, "a")
        self.assertEqual(attr.doc_string, "")
        attr = helper.make_attribute("a", "value", "doc")
        self.assertEqual(attr.name, "a")
        self.assertEqual(attr.doc_string, "doc") 

def test_attr_repeated_float(self):  # type: () -> None
        attr = helper.make_attribute("floats", [1.0, 2.0])
        self.assertEqual(attr.name, "floats")
        self.assertEqual(list(attr.floats), [1.0, 2.0])
        checker.check_attribute(attr) 

def test_attr_repeated_str(self):  # type: () -> None
        attr = helper.make_attribute("strings", ["str1", "str2"])
        self.assertEqual(attr.name, "strings")
        self.assertEqual(list(attr.strings), [b"str1", b"str2"])
        checker.check_attribute(attr) 

def test_attr_repeated_graph_proto(self):  # type: () -> None
        graphs = [GraphProto(), GraphProto()]
        graphs[0].name = "a"
        graphs[1].name = "b"
        attr = helper.make_attribute("graphs", graphs)
        self.assertEqual(attr.name, "graphs")
        self.assertEqual(list(attr.graphs), graphs)
        checker.check_attribute(attr) 

def test_node_with_arg(self):  # type: () -> None
        self.assertTrue(defs.has("Relu"))
        # Note: Relu actually does not need an arg, but let's
        # test it.
        node_def = helper.make_node(
            "Relu", ["X"], ["Y"],
            arg_value=1)
        self.assertEqual(node_def.op_type, "Relu")
        self.assertEqual(list(node_def.input), ["X"])
        self.assertEqual(list(node_def.output), ["Y"])
        self.assertEqual(len(node_def.attribute), 1)
        self.assertEqual(
            node_def.attribute[0],
            helper.make_attribute("arg_value", 1)) 

def add_attribute(node_proto, name, value):
    """Add a new attribute into the node proto."""
    node_proto.attribute.extend([make_attribute(name, value)]) 

def make_onnx(self):
        ''' Generate ONNX model from current graph '''
        self.clean_init()
        nodes = []
        for node in self.node:
            n = NodeProto()
            n.input.extend(node.input)
            n.output.extend(node.output)
            n.name = node.name
            n.op_type = node.op_type
            n.attribute.extend(
                helper.make_attribute(key, value)
                for key, value in sorted(node.attrs.items())
                )
            nodes.append(n)
        
        inputs = []
        initializer = []
        for k,v in self.tensor_dict.items():
            init = numpy_helper.from_array(v, name=k)
            initializer.append(init)
            value_info = helper.make_tensor_value_info(
                name=k,
                elem_type=mapping.NP_TYPE_TO_TENSOR_TYPE[v.dtype],
                shape=list(v.shape)
            )
            inputs.append(value_info)
        
        graph_ = helper.make_graph(
            nodes=nodes,
            name='dlpy_graph',
            inputs=inputs+self.uninitialized,
            outputs=self.output,
            initializer=initializer
        )

        model = helper.make_model(graph_)
        return model 

def verify_pooling(x_shape, kernel_shape, strides, pads, out_shape, mode, auto_pad="NOTSET"):
    x_np = np.random.uniform(size=x_shape).astype('float32')

    if mode == 'max':
        node_type = "MaxPool"
    elif mode == 'average':
        node_type = "AveragePool"
    else:
        raise ValueError("Pool method {} is not supported.".format(mode))

    pool_node = helper.make_node(
        node_type, inputs=["x"], outputs=["y"], kernel_shape=kernel_shape, strides=strides)

    if pads is None:
        pad_attr = helper.make_attribute('auto_pad', auto_pad)
    else:
        pad_attr = helper.make_attribute('pads', pads)
    pool_node.attribute.append(pad_attr)

    if mode == 'max':
        storage_attr = helper.make_attribute('storage_order', 0)
        pool_node.attribute.append(storage_attr)

    graph = helper.make_graph([pool_node],
                              "pooling_test",
                              inputs=[helper.make_tensor_value_info("x",
                                                                    TensorProto.FLOAT, list(x_shape))],
                              outputs=[helper.make_tensor_value_info("y",
                                                                     TensorProto.FLOAT, list(out_shape))])

    model = helper.make_model(graph, producer_name='pooling_test')

    for target, ctx in ctx_list():
        onnx_out = get_onnxruntime_output(model, x_np, 'float32')
        tvm_out = get_tvm_output(
            model, [x_np], target, ctx, out_shape)
        tvm.testing.assert_allclose(onnx_out, tvm_out, rtol=1e-5, atol=1e-5) 

def generate(self):
        updated = False
        if self.attributes:
            updated = True
        elif len([k for k, v in self.input.items() if k != v]) > 0:
            updated = True
        elif len([k for k, v in self.output.items() if k != v]) > 0:
            updated = True

        if not updated:
            return [self.origin]
        else:
            onode = onnx_proto.NodeProto()
            onode.name = self.origin.name
            onode.op_type = self.origin.op_type
            onode.input.extend([self.input.get(i_, i_) for i_ in self.origin.input])
            for input_ in self.initializers:
                if input_.name not in onode.input:
                    onode.input.append(input_.name)
            onode.output.extend([self.output.get(o_, o_) for o_ in self.origin.output])
            onode.doc_string = self.origin.doc_string
            onode.domain = self.origin.domain
            onode.attribute.extend(
                attr for attr in self.origin.attribute if attr.name not in self.attributes)
            onode.attribute.extend(
                helper.make_attribute(attr, self.attributes[attr]) for attr in self.attributes)

            return [onode] 

def const_folding_optimizer(graph, outer_graph=None):
    # type: (onnx.GraphProto, onnx.GraphProto)->onnx.GraphProto
    nodelist, reserved_names = reserve_node_for_embedded_graph(graph.node)
    opt_graph = OnnxGraphContext(graph, nodelist)
    node_status = {}
    for ts_ in graph.output:
        _dfs_calc(opt_graph, opt_graph.tensor_to_node[ts_.name], reserved_names, node_status)

    graph.initializer.extend([numpy_helper.from_array(ts_, nm_) for nm_, ts_ in opt_graph.variables.items()])
    new_nodes = [nd_ for nd_ in nodelist if nd_.name in node_status]
    new_nodes = [nd_ for nd_ in new_nodes if nd_.output[0] not in opt_graph.variables]

    def node_key(nd_):
        return abs(node_status[nd_.name])

    new_nodes.sort(key=node_key)
    pruned_initilizers = _remove_unused_initializers(new_nodes, graph.initializer, reserved_names,
                                                     None if outer_graph is None else outer_graph.initializer)
    del graph.node[:]
    graph.node.extend(new_nodes)
    del graph.initializer[:]
    graph.initializer.extend(pruned_initilizers)

    for nd_ in graph.node:
        for aname_, subgraph_ in OnnxGraphContext.get_attr_graph(nd_).items():
            opt_inner_graph = const_folding_optimizer(subgraph_, graph)
            lst_attrs = list(nd_.attribute)
            del nd_.attribute[:]
            lst_attrs = [helper.make_attribute(aname_, opt_inner_graph) if
                         attr.name == aname_ else attr for attr in lst_attrs]
            nd_.attribute.extend(lst_attrs)

    return graph