Класс NAdam

Класс NAdam реализует оптимизатор Nesterov Adam, который сочетает в себе адаптивный метод Adam с импульсом Нестерова для более быстрой сходимости. Оптимизатор принимает параметры модели и позволяет настраивать скорость обучения, коэффициенты затухания моментов и другие гиперпараметры.

Синтаксис

torch.optim.NAdam(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, momentum_decay=0.004, decoupled_weight_decay=False)

Основные параметры:

params - итерируемый объект с параметрами модели или группами параметров для оптимизации
lr - скорость обучения (learning rate), по умолчанию 0.002
betas - коэффициенты затухания для первого и второго моментов, по умолчанию (0.9, 0.999)
eps - малая константа для численной стабильности, по умолчанию 1e-8
weight_decay - коэффициент регуляризации весов, по умолчанию 0
momentum_decay - коэффициент затухания импульса, по умолчанию 0.004
decoupled_weight_decay - использовать ли разделенную регуляризацию весов, по умолчанию False

Пример

Создадим простую линейную модель и применим к ней оптимизатор NAdam:

import torch
import torch.nn as nn

model = nn.Linear(10, 1)
optimizer = torch.optim.NAdam(model.parameters(), lr=0.001)

print(f"Optimizer: {optimizer.__class__.__name__}")
print(f"Learning rate: {optimizer.param_groups[0]['lr']}")

Результат выполнения кода:

Optimizer: NAdam
Learning rate: 0.001

Пример

Используем NAdam для обучения модели на синтетических данных:

import torch
import torch.nn as nn

torch.manual_seed(0)

model = nn.Linear(1, 1)
criterion = nn.MSELoss()
optimizer = torch.optim.NAdam(model.parameters(), lr=0.01)

x = torch.linspace(-1, 1, 100).reshape(-1, 1)
y = 2 * x + 1 + 0.1 * torch.randn_like(x)

for epoch in range(50):
	optimizer.zero_grad()
	output = model(x)
	loss = criterion(output, y)
	loss.backward()
	optimizer.step()
	
	if epoch % 20 == 0:
		print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

Результат выполнения кода:

Epoch 0, Loss: 1.1891
Epoch 20, Loss: 0.0106
Epoch 40, Loss: 0.0099

Пример

Настройка параметров оптимизатора для более сложной модели:

import torch
import torch.nn as nn

torch.manual_seed(0)

model = nn.Sequential(
	nn.Linear(20, 64),
	nn.ReLU(),
	nn.Linear(64, 1)
)

optimizer = torch.optim.NAdam(
	model.parameters(),
	lr=0.001,
	betas=(0.9, 0.999),
	momentum_decay=0.004,
	weight_decay=1e-5
)

print("Model architecture:")
print(model)
print(f"Optimizer parameters: {optimizer.defaults}")

Результат выполнения кода:

Model architecture:
Sequential(
  (0): Linear(in_features=20, out_features=64, bias=True)
  (1): ReLU()
  (2): Linear(in_features=64, out_features=1, bias=True)
)
Optimizer parameters: {'lr': 0.001, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 1e-05, 'amsgrad': False, 'momentum_decay': 0.004, 'decoupled_weight_decay': False}

Смотрите также

класс Adam,
который реализует стандартный адаптивный оптимизатор
класс AdamW,
который добавляет разделенную регуляризацию весов
класс SGD,
который реализует стохастический градиентный спуск с импульсом
класс RMSprop,
который использует среднеквадратичное распространение градиента