megengine.traced_module.trace_module

trace_module(mod, *args, **kwargs)[source]

Traces module mod and returns corresponding TracedModule.

Parameters

  • mod (module.Module) – the module will be converted to TracedModule.

  • args (Tuple[Any]]) – the positional arguments passed to forward method of mod.

  • kwargs (Dict[str, Any]) – the keyword arguments passed to forward method of mod.

Returns

TracedModule. The TracedModule object convert from input Module mod.

Return type

Return type

Examples

import megengine.functional as F
import megengine.module as M
import megengine as mge
from model import resnet18

# resnet : Module
resnet = resnet18()

import megengine.traced_module as tm
inp = F.zeros(shape=(1,3,224,224))

# traced_resnet : TracedModule
traced_resnet =  tm.trace_module(resnet, inp)

previous

TracedModule

next

megengine.traced_module.register_as_builtin