megengine.traced_module.traced_module.InternalGraph.reset_outputs

InternalGraph.reset_outputs(outputs)[源代码]

Reset the output Nodes of the graph.

注解

This method only supports resetting the output of graphs that do not have a parent graph.

参数

outputs – an object which inner element is Node. Support tuple, list dict, etc.

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.reset_outputs((out_node, {"input": inp_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
}
(Tensor([1.], device=xpux:0), {'input': Tensor([0.], device=xpux:0)})