torch.nn.utils.parametrize.cached

torch.nn.utils.parametrize.cached() [исходный код]

Менеджер контекста, включающий систему кэширования для параметризаций, зарегистрированных с помощью register_parametrization().

Когда этот менеджер контекста активен, значение параметризованных объектов вычисляется и кэшируется при первом обращении к ним. При выходе из менеджера контекста кэшированные значения удаляются.

Это полезно, когда параметризованный параметр используется более одного раза во время прямого прохода. Например, при параметризации рекуррентного ядра RNN или при совместном использовании весов.

Самый простой способ активировать кэш — обернуть прямой проход нейронной сети

import torch.nn.utils.parametrize as P

...
with P.cached():
    output = model(inputs)

при обучении и оценке. Также можно обернуть части модулей, в которых параметризованные тензоры используются несколько раз. Например, цикл RNN с параметризованным рекуррентным ядром:

with P.cached():
    for x in xs:
        out_rnn = self.rnn_cell(x, out_rnn)