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
andout
: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)})