megengine.traced_module.traced_module.InternalGraph.add_output_node¶

InternalGraph.add_output_node(node)[源代码]¶

Add an output node to the Graph.

The Graph output will become a tuple after calling add_output_node. The first element of the tuple is the original output, and the second is the node.

For example, the following code

import megengine.functional as F
import megengine.module as M
import megengine.traced_module as tm

class MyModule(M.Module):
    def forward(self, x):
        x = x + 1
        return x

net = MyModule()

inp = F.zeros(shape = (1, ))
traced_module = tm.trace_module(net, inp)
graph = traced_module.graph
inp_node = graph.inputs[1]
out_node = graph.outputs[0]
graph.add_output_node(inp_node)
graph.add_output_node(out_node)
out = traced_module(inp)

Will produce the following InternalGraph and out:

print(graph)
print(out)

MyModule.Graph (self, x) {
        %2:     add_out = x.__add__(1, )
        return add_out, x, add_out
}
((Tensor([1.], device=xpux:0), Tensor([0.], device=xpux:0)), Tensor([1.], device=xpux:0))

megengine.traced_module.traced_module.InternalGraph.add_input_node

megengine.traced_module.traced_module.InternalGraph.compile