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]

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