Python PyTorch ScriptModule.apply用法及代碼示例

本文簡要介紹python語言中 torch.jit.ScriptModule.apply 的用法。

用法:

apply(fn)

參數：

fn(Module-> 無)-要應用於每個子模塊的函數

返回：

self

返回類型：

torch.nn.Module

將 fn 遞歸地應用於每個子模塊(由 .children() 返回)以及自身。典型用途包括初始化模型的參數(另見 torch.nn.init)。

例子：

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)