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)
)