torch.optim.swa_utils.get_ema_avg_fn

torch.optim.swa_utils.get_ema_avg_fn(decay=0.999) [исходный код]

Возвращает функцию, применяющую экспоненциальное скользящее среднее (EMA) для нескольких параметров.

EMA вычисляется следующим образом:

$$W_0^{\text{EMA}} = W_0^{\text{model}}$$

$$W_{t+1}^{\text{EMA}} = \text{decay} \times W_t^{\text{EMA}} + (1 - \text{decay}) \times W_{t+1}^{\text{model}}$$

где $W_t^{\text{EMA}}$ — параметр EMA на шаге $t$, $W_t^{\text{model}}$ — параметр модели на шаге $t$, а $\text{decay}$ — коэффициент затухания (по умолчанию: 0.999).

Параметры:

decay (float) – Коэффициент затухания для EMA. Должен находиться в диапазоне [0, 1]. По умолчанию: 0.999

Возвращает:

Функцию, которая обновляет параметры EMA с учетом текущих параметров модели

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

Callable