LazyModuleMixin

class torch.nn.modules.lazy.LazyModuleMixin(*args, **kwargs) [исходный код]

Примесь для модулей, которые лениво инициализируют параметры; такие модули также называются «ленивыми модулями».

Предупреждение

Ленивые модули — это новая экспериментальная функция, которая активно разрабатывается, и её API, вероятно, изменится.

Модули, которые лениво инициализируют параметры, или «ленивые модули», определяют формы своих параметров по первому входному значению (или первым входным значениям) метода forward. До первого вызова forward они содержат torch.nn.UninitializedParameter, к которым не следует обращаться или которые не следует использовать, а после него содержат обычные torch.nn.Parameter. Ленивые модули удобны, поскольку позволяют не вычислять некоторые аргументы модуля, например аргумент in_features типичного torch.nn.Linear.

После создания сети с ленивыми модулями сначала следует преобразовать её к нужному типу данных и разместить на требуемом устройстве. Это связано с тем, что ленивые модули только определяют формы, поэтому применяется обычное поведение при выборе типа данных и размещении на устройстве. Затем ленивые модули должны выполнить «пробные запуски», чтобы инициализировать все компоненты модуля. Во время этих «пробных запусков» через сеть и каждый из её ленивых модулей передаются входные данные нужного размера, типа и устройства. После этого сеть можно использовать как обычно.

>>> class LazyMLP(torch.nn.Module):
...     def __init__(self) -> None:
...         super().__init__()
...         self.fc1 = torch.nn.LazyLinear(10)
...         self.relu1 = torch.nn.ReLU()
...         self.fc2 = torch.nn.LazyLinear(1)
...         self.relu2 = torch.nn.ReLU()
...
...     def forward(self, input):
...         x = self.relu1(self.fc1(input))
...         y = self.relu2(self.fc2(x))
...         return y
>>> # constructs a network with lazy modules
>>> lazy_mlp = LazyMLP()
>>> # transforms the network's device and dtype
>>> # NOTE: these transforms can and should be applied after construction and before any 'dry runs'
>>> lazy_mlp = lazy_mlp.cuda()
>>> lazy_mlp
LazyMLP( (fc1): LazyLinear(in_features=0, out_features=10, bias=True)
  (relu1): ReLU()
  (fc2): LazyLinear(in_features=0, out_features=1, bias=True)
  (relu2): ReLU()
)
>>> # performs a dry run to initialize the network's lazy modules
>>> lazy_mlp(torch.ones(10, 10).cuda())
>>> # after initialization, LazyLinear modules become regular Linear modules
>>> lazy_mlp
LazyMLP(
  (fc1): Linear(in_features=10, out_features=10, bias=True)
  (relu1): ReLU()
  (fc2): Linear(in_features=10, out_features=1, bias=True)
  (relu2): ReLU()
)
>>> # attaches an optimizer, since parameters can now be used as usual
>>> optim = torch.optim.SGD(lazy_mlp.parameters(), lr=0.01)

При использовании ленивых модулей следует также учитывать, что порядок инициализации параметров сети может измениться, поскольку ленивые модули всегда инициализируются после остальных модулей. Например, если бы класс LazyMLP, определённый выше, сначала содержал модуль torch.nn.LazyLinear, а затем обычный модуль torch.nn.Linear, второй модуль был бы инициализирован при создании, а первый — во время первого пробного запуска. Из-за этого параметры сети с ленивыми модулями могут инициализироваться иначе, чем параметры сети без ленивых модулей: порядок инициализации параметров, который часто зависит от генератора случайных чисел с состоянием, различается. Подробнее см. в разделе Воспроизводимость.

Ленивые модули, как и другие модули, можно сериализовать с помощью словаря состояния. Например:

>>> lazy_mlp = LazyMLP()
>>> # The state dict shows the uninitialized parameters
>>> lazy_mlp.state_dict()
OrderedDict({'fc1.weight': <UninitializedParameter>,
             'fc1.bias': <UninitializedParameter>,
             'fc2.weight': <UninitializedParameter>,
             'fc2.bias': <UninitializedParameter>})

Ленивые модули могут загружать обычные torch.nn.Parameter (то есть инициализированные LazyModules можно сериализовать и десериализовать, и они останутся инициализированными).

>>> full_mlp = LazyMLP()
>>> # Dry run to initialize another module
>>> full_mlp.forward(torch.ones(10, 1))
>>> # Load an initialized state into a lazy module
>>> lazy_mlp.load_state_dict(full_mlp.state_dict())
>>> # The state dict now holds valid values
>>> lazy_mlp.state_dict()
OrderedDict([('fc1.weight',
              tensor([[-0.3837],
                      [ 0.0907],
                      [ 0.6708],
                      [-0.5223],
                      [-0.9028],
                      [ 0.2851],
                      [-0.4537],
                      [ 0.6813],
                      [ 0.5766],
                      [-0.8678]])),
             ('fc1.bias',
              tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,
                       4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),
             ('fc2.weight',
              tensor([[ 0.1320,  0.2938,  0.0679,  0.2793,  0.1088, -0.1795, -0.2301,  0.2807,
                        0.2479,  0.1091]])),
             ('fc2.bias', tensor([0.0019]))])

Однако обратите внимание: загруженные параметры не будут заменены во время «пробного запуска», если они инициализированы при загрузке состояния. Это не позволяет использовать инициализированные модули в разных контекстах.

has_uninitialized_params() [исходный код]

Проверяет, есть ли в модуле неинициализированные параметры.

initialize_parameters(*args, **kwargs) [исходный код]

Инициализирует параметры в соответствии со свойствами входного пакета данных.

Этот метод позволяет отделить инициализацию параметров от прямого прохода при определении форм параметров.