Adafactor

class torch.optim.Adafactor(params, lr=0.01, beta2_decay=-0.8, eps=(None, 0.001), d=1.0, weight_decay=0.0, *, foreach=None, maximize=False)

Реализует алгоритм Adafactor.

$$\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{(lr)}, \: \tau \text{(}\beta_2\text{ decay)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)}, \\ &\hspace{15mm} \: \epsilon_1, \epsilon_2 \text{ (epsilons)}, \: d \text{(clipping threshold)}, \\ &\hspace{15mm} \: \lambda \text{(weight decay)}, \: \textit{maximize} \\ &\textbf{initialize} : \: R_0 \leftarrow 0 \text{ (second moment row factor)}, \\ &\hspace{23mm} \: C_0 \leftarrow 0 \text{ (second moment col factor)}, \\ &\hspace{23mm} \: \widehat{V}_0 \leftarrow 0 \text{ (second moment for vectors)} \\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ &\hspace{10mm}G_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}G_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ &\hspace{5mm}\widehat{\beta}_{2_t} \leftarrow 1 - t^{\tau} \\ &\hspace{5mm}\rho_t \leftarrow min(\gamma, \frac{1}{\sqrt{t}}) \\ &\hspace{5mm}\alpha_t \leftarrow max(\epsilon_2, \text{RMS}(\theta_{t-1}))\rho_t \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ &\hspace{5mm}\textbf{if} \: \text{dim}(G_t) > 1: \\ &\hspace{10mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t) \cdot 1_m \\ &\hspace{10mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t) \\ &\hspace{10mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}\widehat{V}_t \leftarrow \widehat{\beta}_{2_t}\widehat{V}_{t-1}+ (1-\widehat{\beta}_{2_t}) \cdot (G_t \odot G_t) \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ &\hspace{5mm}\widehat{U}_t \leftarrow \frac{U_t}{max(1, \frac{\text{RMS}(U_t)}{d})} \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \alpha_t \widehat{U}_t \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned}$$

Дополнительные сведения об алгоритме см. в статье Adafactor: Adaptive Learning Rates with Sublinear Memory Cost.

Параметры:

• params (iterable) – итерируемый объект с параметрами или named_parameters для оптимизации либо итерируемый объект словарей, определяющих группы параметров. При использовании named_parameters все параметры во всех группах должны иметь имена
• lr (float, Tensor, необязательный) – в отличие от других оптимизаторов, Adafactor не требует скорости обучения, а Noam Shazeer и Mitchell Stern вообще не используют lr. В отступление от статьи, эта реализация использует lr для применения затухания весов и в качестве максимального значения относительного размера шага rho_t. Обратите внимание: в статье в качестве максимального значения относительного размера шага используется константа 0.01, поэтому мы задаём 0.01 в качестве значения по умолчанию. (по умолчанию: 1e-2)
• beta2_decay (float, необязательный) – скорость затухания beta2. Обычно beta2 обозначает коэффициент, используемый для вычисления скользящего среднего квадрата градиента. (по умолчанию: -0.8)
• eps (Tuple[float, float], необязательный) – epsilon1 — слагаемое, добавляемое к знаменателю при вычислении обновления для повышения численной устойчивости. Такое использование epsilon1 отличается от алгоритма, описанного в статье! Подробнее см. примечание ниже. epsilon2 используется для предотвращения слишком малого обновления весов при масштабировании параметров. (по умолчанию: (None, 1e-3))
• d (float, необязательный) – порог отсечения, используемый для предотвращения чрезмерно больших обновлений.
• weight_decay (float, необязательный) – коэффициент затухания весов (по умолчанию: 1e-2)
• foreach (bool, необязательный) – использовать ли реализацию оптимизатора с foreach. Обратите внимание: реализация с foreach использует примерно на sizeof(params) больше пиковой памяти, чем версия с циклом for, поскольку промежуточные данные представлены списком тензоров, а не одним тензором. Поскольку Adafactor часто используют в условиях ограниченной памяти, по умолчанию применяется более медленная реализация с циклом for и одним тензором, если этот флаг явно не установлен в True. Такое поведение отличается от поведения других оптимизаторов, которые для ускорения работы пытаются по умолчанию использовать foreach на CUDA. (по умолчанию: None)
• maximize (bool, необязательный) – максимизировать целевую функцию по параметрам вместо минимизации (по умолчанию: False)

Примечание

Реализация Adafactor имеет небольшие отличия от подхода Noam Shazeer и Mitchell Stern и реализаций в некоторых других платформах: они касаются использования скорости обучения и $\epsilon_1$.

Что касается гиперпараметра скорости обучения: Noam Shazeer и Mitchell Stern вообще не используют lr, поскольку в описанном алгоритме для управления размером шага используются $\rho_t$ и отсечение обновлений.

В этой реализации lr влияет на максимальное значение $\rho_t$:

$$\begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(\gamma, \frac{1}{\sqrt{t}}) \end{aligned}$$

Это отличается от подхода Noam Shazeer и Mitchell Stern, которые используют константу 0.01 в качестве максимального значения $\rho_t$

$$\begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(0.01, \frac{1}{\sqrt{t}}) \end{aligned}$$

Noam Shazeer и Mitchell Stern не предписывают конкретный способ вычисления затухания весов, поэтому мы используем скорость обучения в качестве коэффициента для отделённого от градиента затухания весов, как предлагается в статье Decoupled Weight Decay Regularization.

Что касается использования $\epsilon_1$: в реализации предпринята попытка воспроизвести предполагаемое намерение Noam Shazeer и Mitchell Stern — использовать $\epsilon_1$ в качестве стабилизирующего слагаемого, когда квадрат градиента становится малым.

Эту стабилизацию можно записать следующим образом:

$$\begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ \end{aligned}$$

При этом множители строк и столбцов квадрата градиента $R_t$ и $C_t$ остаются без изменений, а $\epsilon_1$ применяется при окончательном вычислении оценки дисперсии $\widehat{V}_t$ и при вычислении обновления $U_t$.

Это отличается от подхода Noam Shazeer и Mitchell Stern и других платформ, в которых $\epsilon_1$ добавляют к множителям строк и столбцов квадрата градиента, но не используют в последующих вычислениях:

$$\begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{1^\top_n \cdot R_t} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{\sqrt{\widehat{V}_t}} \\ \end{aligned}$$

Обратите внимание: Noam Shazeer и Mitchell Stern описывают использование суммы квадратов градиентов, тогда как в этой реализации используется среднее. Математически эти варианты эквивалентны, а использование среднего обеспечивает большую численную устойчивость при больших суммах.

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

Добавляет группу параметров в param_groups Optimizer.

Это может быть полезно при дообучении предварительно обученной сети: замороженные слои можно сделать обучаемыми и добавлять в Optimizer по мере обучения.

Параметры:

param_group (dict) – задаёт, какие тензоры следует оптимизировать, а также параметры оптимизации, специфичные для группы.

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

Загрузить состояние оптимизатора.

Параметры:

state_dict (dict) – состояние оптимизатора. Должен быть объектом, возвращённым вызовом state_dict().

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

Убедитесь, что этот метод вызывается после инициализации torch.optim.lr_scheduler.LRScheduler, поскольку его вызов до этого перезапишет загруженные скорости обучения.

Примечание

Имена параметров (если они указаны под ключом «param_names» каждой группы параметров в state_dict()) не влияют на процесс загрузки. Чтобы использовать имена параметров в пользовательских случаях (например, когда параметры в загруженном словаре состояния отличаются от инициализированных в оптимизаторе), необходимо реализовать пользовательский register_load_state_dict_pre_hook для соответствующей адаптации загруженного словаря. Если param_names существуют в загруженном словаре состояния param_groups, они будут сохранены и переопределят текущие имена, если они есть, в состоянии оптимизатора. Если их нет в загруженном словаре состояния, param_names оптимизатора останутся без изменений.

Пример

>>> optimizer = ...  # initialized optimizer matching the saved state
>>> scheduler1 = torch.optim.lr_scheduler.LinearLR(
...     optimizer,
...     start_factor=0.1,
...     end_factor=1,
...     total_iters=20,
... )
>>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
...     optimizer,
...     T_max=80,
...     eta_min=3e-5,
... )
>>> lr = torch.optim.lr_scheduler.SequentialLR(
...     optimizer,
...     schedulers=[scheduler1, scheduler2],
...     milestones=[20],
... )
>>> lr.load_state_dict(torch.load("./save_seq.pt"))
>>> # now load the optimizer checkpoint after loading the LRScheduler
>>> optimizer.load_state_dict(torch.load("./save_optim.pt"))

register_load_state_dict_post_hook(hook, prepend=False) [исходный код]

Зарегистрировать пост-хук load_state_dict, который будет вызван после вызова load_state_dict(). Он должен иметь следующую сигнатуру:

hook(optimizer) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора.

Хук будет вызван с аргументом self после вызова load_state_dict для self. Зарегистрированный хук можно использовать для постобработки после того, как load_state_dict загрузил state_dict.

Параметры:

• hook (Callable) – Пользовательский хук для регистрации.
• prepend (bool) – Если значение True, предоставленный пост-hook будет вызван перед всеми уже зарегистрированными пост-хуками для load_state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных пост-хуков. (по умолчанию: False)

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook, prepend=False) [исходный код]

Зарегистрировать пре-хук load_state_dict, который будет вызван до вызова load_state_dict(). Он должен иметь следующую сигнатуру:

hook(optimizer, state_dict) -> state_dict or None

Аргумент optimizer — это используемый экземпляр оптимизатора, а аргумент state_dict — поверхностная копия state_dict, переданного пользователем в load_state_dict. Хук может изменить state_dict на месте или при необходимости вернуть новый. Если возвращён state_dict, он будет загружен в оптимизатор.

Хук будет вызван с аргументами self и state_dict до вызова load_state_dict для self. Зарегистрированный хук можно использовать для предварительной обработки перед вызовом load_state_dict.

Параметры:

• hook (Callable) – Пользовательский хук для регистрации.
• prepend (bool) – Если значение True, предоставленный пре-hook будет вызван перед всеми уже зарегистрированными пре-хуками для load_state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных пре-хуков. (по умолчанию: False)

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_state_dict_post_hook(hook, prepend=False) [исходный код]

Зарегистрировать пост-хук словаря состояния, который будет вызван после вызова state_dict().

Он должен иметь следующую сигнатуру:

hook(optimizer, state_dict) -> state_dict or None

Хук будет вызван с аргументами self и state_dict после создания state_dict для self. Хук может изменить state_dict на месте или при необходимости вернуть новый. Зарегистрированный хук можно использовать для постобработки state_dict перед его возвратом.

Параметры:

• hook (Callable) – Пользовательский хук для регистрации.
• prepend (bool) – Если значение True, предоставленный пост-hook будет вызван перед всеми уже зарегистрированными пост-хуками для state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных пост-хуков. (по умолчанию: False)

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

register_state_dict_pre_hook(hook, prepend=False) [исходный код]

Зарегистрировать пре-хук словаря состояния, который будет вызван до вызова state_dict().

Он должен иметь следующую сигнатуру:

hook(optimizer) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора. Хук будет вызван с аргументом self до вызова state_dict для self. Зарегистрированный хук можно использовать для предварительной обработки перед вызовом state_dict.

Параметры:

• hook (Callable) – Пользовательский хук для регистрации.
• prepend (bool) – Если значение True, предоставленный пре-hook будет вызван перед всеми уже зарегистрированными пре-хуками для state_dict. В противном случае предоставленный hook будет вызван после всех уже зарегистрированных пре-хуков. (по умолчанию: False)

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

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

Зарегистрировать пост-хук шага оптимизатора, который будет вызван после шага оптимизатора.

Он должен иметь следующую сигнатуру:

hook(optimizer, args, kwargs) -> None

Аргумент optimizer — это используемый экземпляр оптимизатора.

Параметры:

hook (Callable) – Пользовательский хук для регистрации.

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

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

Зарегистрировать пре-хук шага оптимизатора, который будет вызван перед шагом оптимизатора.

Он должен иметь следующую сигнатуру:

hook(optimizer, args, kwargs) -> None or modified args and kwargs

Аргумент optimizer — это используемый экземпляр оптимизатора. Если пре-хук изменяет args и kwargs, преобразованные значения возвращаются в виде кортежа, содержащего new_args и new_kwargs.

Параметры:

hook (Callable) – Пользовательский хук для регистрации.

Возвращает:

объект-дескриптор, который можно использовать для удаления добавленного хука вызовом handle.remove()

Тип возвращаемого значения:

torch.utils.hooks.RemovableHandle

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

Вернуть состояние оптимизатора в виде dict.

Он содержит две записи:

• state: a Dict holding current optimization state. Its content

  различается у разных классов оптимизаторов, но имеет ряд общих характеристик. Например, состояние сохраняется для каждого параметра, а сам параметр НЕ сохраняется. state — это словарь, сопоставляющий идентификаторы параметров со словарём состояния соответствующего параметра.

• param_groups: a List containing all parameter groups where each

  Группа параметров представляет собой словарь. Каждая группа параметров содержит метаданные, специфичные для оптимизатора, например скорость обучения и затухание весов, а также список идентификаторов параметров, входящих в группу. Если группа параметров была инициализирована с помощью named_parameters(), содержимое names также будет сохранено в словаре состояния.

ПРИМЕЧАНИЕ: Идентификаторы параметров могут выглядеть как индексы, но это просто идентификаторы, связывающие состояние с param_group. При загрузке из state_dict оптимизатор сопоставит params группы параметров (целочисленные идентификаторы) и param_groups оптимизатора (фактические nn.Parameter s), чтобы сопоставить состояния БЕЗ дополнительной проверки.

Возвращаемый словарь состояния может выглядеть примерно так:

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}

Тип возвращаемого значения:

dict[str, Any]

step(closure=None) [исходный код]

Выполнить один шаг оптимизации.

Параметры:

closure (Callable, optional) – Замыкание, которое повторно вычисляет модель и возвращает значение функции потерь.

zero_grad(set_to_none=True) [исходный код]

Сбросить градиенты всех оптимизируемых torch.Tensor s.

Параметры:

set_to_none (bool, optional) –

Вместо установки нулевого значения установить для градиентов значение None. По умолчанию: True

В целом это уменьшает использование памяти и может немного повысить производительность. Однако при этом меняется поведение в некоторых случаях. Например:

1. Если пользователь пытается получить доступ к градиенту и выполнить над ним операции вручную, атрибут None и Tensor, заполненный нулями, будут вести себя по-разному.
2. Если пользователь запрашивает zero_grad(set_to_none=True), а затем выполняет обратный проход, для .grad, не получивших градиент, гарантированно будет установлено значение None.
3. torch.optim оптимизаторы ведут себя по-разному, если градиент равен 0 или None (в одном случае выполняется шаг с градиентом 0, а в другом шаг полностью пропускается).