LazyModuleMixin

class torch.nn.modules.lazy.LazyModuleMixin(*args, **kwargs) [source]

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

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

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

>>> class LazyMLP(torch.nn.Module):
...    def __init__(self):
...        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().double()
>>> 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(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', Uninitialized parameter),
             ('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', Uninitialized parameter),
             ('fc2.bias', tensor([0.0019]))])

Ленивые модули могут загружать обычные 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() [source]

Проверка наличия неинициализированных параметров в модуле.

initialize_parameters(*args, **kwargs) [source]

Инициализация параметров в соответствии с свойствами входного набора данных. Это добавляет интерфейс для изоляции инициализации параметров от прохода forward при вычислении формы параметров.