當前位置: 首頁>>代碼示例 >>用法及示例精選 >>正文


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

相關用法


注:本文由純淨天空篩選整理自pytorch.org大神的英文原創作品 torchrec.modules.lazy_extension.LazyModuleExtensionMixin.apply。非經特殊聲明,原始代碼版權歸原作者所有,本譯文未經允許或授權,請勿轉載或複製。