Python PyTorch LazyModuleExtensionMixin.apply用法及代码示例

本文简要介绍python语言中 torchrec.modules.lazy_extension.LazyModuleExtensionMixin.apply 的用法。

用法:

apply(fn: Callable[[torch.nn.modules.module.Module], None]) → torch.nn.modules.module.Module

参数：

fn(torch.nn.Module -> 无) -要应用于每个子模块的函数。

返回：

self

返回类型：

火炬.nn.模块

将 fn 递归地应用于每个子模块(由 .children() 返回)以及自身。典型用途包括初始化模型的参数。

注意

在未初始化的 lazy-module 上调用 apply() 将导致错误。在调用lazy-module 上的apply() 之前，用户需要初始化lazy-module(通过执行虚拟前向传递)。

例子：

@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == torch.nn.LazyLinear:
        m.weight.fill_(1.0)
        print(m.weight)

linear = torch.nn.LazyLinear(2)
linear.apply(init_weights)  # this fails, because `linear` (a lazy-module) hasn't been initialized yet

input = torch.randn(2, 10)
linear(input)  # run a dummy forward pass to initialize the lazy-module

linear.apply(init_weights)  # this works now