Adam

class torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False, *, foreach=None, maximize=False, capturable=False, differentiable=False, fused=None, decoupled_weight_decay=False) [исходный код]

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

$$\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)} \\ &\hspace{13mm} \lambda \text{ (weight decay)}, \: \textit{amsgrad}, \:\textit{maximize}, \: \epsilon \text{ (epsilon)} \\ &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, v_0\leftarrow 0 \text{ (second moment)},\: v_0^{max}\leftarrow 0 \\[-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}\textbf{if} \: \lambda \neq 0 \\ &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ &\hspace{5mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ &\hspace{5mm}\textbf{if} \: amsgrad \\ &\hspace{10mm} v_t^{max} \leftarrow \mathrm{max}(v_{t-1}^{max},v_t) \\ &\hspace{10mm}\widehat{v_t} \leftarrow v_t^{max}/\big(1-\beta_2^t \big) \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned}$$

Дополнительные сведения об алгоритме см. в статье Adam: метод стохастической оптимизации.

Параметры:

• params (итерируемый объект) – итерируемый объект, содержащий параметры или именованные параметры для оптимизации, либо итерируемый объект словарей, определяющих группы параметров. При использовании именованных параметров все параметры во всех группах должны иметь имена
• lr (float, Tensor, необязательно) – скорость обучения (по умолчанию: 1e-3). Поддержка тензорного значения LR пока реализована не во всех вариантах. Используйте значение LR типа float, если вы также не задаёте fused=True или capturable=True.
• betas (tuple[float | Tensor, float | Tensor], необязательно) – коэффициенты для вычисления скользящих средних градиента и его квадрата. Если передан тензор, он должен содержать один элемент. (по умолчанию: (0.9, 0.999))
• eps (float, необязательно) – слагаемое, добавляемое к знаменателю для повышения численной устойчивости (по умолчанию: 1e-8)
• weight_decay (float, необязательно) – затухание весов (штраф L2) (по умолчанию: 0)
• decoupled_weight_decay (bool, необязательно) – если True, этот оптимизатор эквивалентен AdamW, и алгоритм не будет накапливать затухание весов в импульсе и дисперсии. (по умолчанию: False)
• amsgrad (bool, необязательно) – использовать ли вариант AMSGrad этого алгоритма из статьи О сходимости Adam и дальнейших улучшениях (по умолчанию: False)
• foreach (bool, необязательно) – использовать ли реализацию оптимизатора foreach. Если пользователь не указал значение (то есть foreach равен None), на CUDA будет предпринята попытка использовать foreach вместо реализации с циклом for, поскольку обычно она значительно производительнее. Обратите внимание: из-за промежуточных результатов в виде списка тензоров, а не одного тензора, реализация foreach использует примерно на sizeof(params) больше пиковой памяти, чем вариант с циклом for. Если памяти недостаточно, обрабатывайте за один раз меньше параметров через оптимизатор или задайте этому флагу значение False (по умолчанию: None)
• maximize (bool, необязательно) – максимизировать целевую функцию относительно params вместо её минимизации (по умолчанию: False)
• capturable (bool, необязательно) – безопасно ли захватывать этот экземпляр в граф, например в графы CUDA или для поддержки torch.compile. Тензоры можно захватывать только на поддерживаемых ускорителях. Значение True может снизить производительность вне графа, поэтому, если вы не собираетесь захватывать этот экземпляр в граф, оставьте значение False (по умолчанию: False)
• differentiable (bool, необязательно) – должна ли выполняться автоматическая дифференциация через шаг оптимизатора при обучении. В противном случае функция step() выполняется в контексте torch.no_grad(). Значение True может снизить производительность, поэтому оставьте False, если не планируете выполнять автоматическую дифференциацию через этот экземпляр (по умолчанию: False)
• fused (bool, необязательно) – использовать ли объединённую реализацию. В настоящее время поддерживаются torch.float64, torch.float32, torch.float16 и torch.bfloat16. (по умолчанию: None)

Примечание

Реализации foreach и fused обычно быстрее варианта с циклом for и одним тензором; теоретически самая быстрая — fused, поскольку она использует вертикальное и горизонтальное объединение. Поэтому, если пользователь не указал ни один из флагов (то есть foreach = fused = None), мы попытаемся по умолчанию использовать реализацию foreach, если все тензоры находятся на CUDA. Почему не fused? Поскольку реализация fused относительно новая, мы хотим дать ей достаточно времени на проверку в реальных условиях. Чтобы указать fused, передайте True для fused. Чтобы принудительно использовать реализацию с циклом for, передайте False для foreach или fused.

Примечание

Прототипная реализация Adam и AdamW для MPS поддерживает torch.float32 и torch.float16.

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

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

Это может быть полезно при дообучении предварительно обученной сети: замороженные слои можно сделать обучаемыми и добавлять в 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 — это словарь, сопоставляющий идентификаторы параметров с Dict, содержащим состояние каждого параметра.

• param_groups: a List containing all parameter groups where each

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

ПРИМЕЧАНИЕ: Идентификаторы параметров могут выглядеть как индексы, но это просто идентификаторы, связывающие состояние с 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 и тензор, заполненный нулями, ведут себя по-разному.
2. Если пользователь запрашивает zero_grad(set_to_none=True), а затем выполняет обратное распространение, для параметров, не получивших градиент, .grad гарантированно будут равны None.
3. Оптимизаторы torch.optim ведут себя по-разному в зависимости от того, равен градиент 0 или None (в одном случае выполняется шаг с градиентом 0, а в другом шаг полностью пропускается).