Python PyTorch lazy_apply用法及代碼示例

本文簡要介紹python語言中 torchrec.modules.lazy_extension.lazy_apply 的用法。

用法:

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

參數：

• module(火炬.nn.模塊) -遞歸應用fn 的模塊。

• fn(可調用[[火炬.nn.模塊],None]) - 要附加的函數module然後應用於每個子模塊module和module本身。

返回：

module 附有 fn。

返回類型：

火炬.nn.模塊

將函數附加到模塊，該函數將遞歸地應用於模塊的每個子模塊(由 .children() 返回)以及模塊本身在第一次前向傳遞之後(即在所有子模塊和參數都已初始化之後)。

典型用途包括初始化惰性模塊(即從 LazyModuleMixin 繼承的模塊)的參數的數值。

注意

lazy_apply() 可用於惰性和非惰性模塊。

例子：

@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)
lazy_apply(linear, init_weights)  # doesn't run `init_weights` immediately
input = torch.randn(2, 10)
linear(input)  # runs `init_weights` only once, right after first forward pass

seq = torch.nn.Sequential(torch.nn.LazyLinear(2), torch.nn.LazyLinear(2))
lazy_apply(seq, init_weights)  # doesn't run `init_weights` immediately
input = torch.randn(2, 10)
seq(input)  # runs `init_weights` only once, right after first forward pass