oneflow.nn.Module.apply

Module.apply(fn)

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model.

Parameters

fn (Module -> None) – function to be applied to each submodule

Returns

self

Return type

Module

Example:

>>> import oneflow as flow
>>> import oneflow.nn as nn
>>> @flow.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)
tensor([[1., 1.],
        [1., 1.]], dtype=oneflow.float32, requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
tensor([[1., 1.],
        [1., 1.]], dtype=oneflow.float32, requires_grad=True)
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)
)