megengine.traced_module.trace_module

trace_module(mod, *args, **kwargs)

追溯 mod 模块并返回对应的 TracedModule.

参数

• mod (module.Module) – 这个模块将被转成 TracedModule.

• args (Tuple[Any]]) – 传递给 mod 的 forward 方法的位置参数。

• kwargs (Dict[str, Any]) – 传递给 mod 的 forward 方法的关键字参数。

返回

TracedModule. The TracedModule object convert from input Module mod.

返回类型

Return type

实际案例

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)

上一页

TracedModule

下一页

megengine.traced_module.register_as_builtin