Класс SGD

Класс SGD (Stochastic Gradient Descent) реализует алгоритм стохастического градиентного спуска для оптимизации параметров модели. Первым параметром конструктор принимает список оптимизируемых параметров (обычно model.parameters). Вторым параметром можно передать скорость обучения lr. Также класс поддерживает настройку импульса (momentum), весового затухания (weight_decay) и другие гиперпараметры.

Синтаксис

torch.optim.SGD(
	params,
	lr=0.01,
	momentum=0,
	weight_decay=0,
	nesterov=False
)

Пример

Давайте создадим простую линейную модель и оптимизатор SGD со скоростью обучения 0.1:

import torch

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
print(optimizer)

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

SGD (
Parameter Group 0
	dampening: 0
	foreach: None
	initial_lr: 0.1
	lr: 0.1
	maximize: False
	momentum: 0
	nesterov: False
	weight_decay: 0
)

Пример

Используем оптимизатор SGD для одного шага обучения. Сначала определяем модель и функцию потерь, затем вычисляем градиенты и обновляем веса с помощью метода step:

import torch

torch.manual_seed(0)

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

x = torch.tensor([[1.0, 2.0]])
y = torch.tensor([[1.0]])

criterion = torch.nn.MSELoss()

optimizer.zero_grad()
pred = model(x)
loss = criterion(pred, y)
loss.backward()
optimizer.step()

print(loss.item())

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

0.8130650511474609

Пример

Оптимизатор SGD поддерживает импульс (momentum), который помогает ускорить сходимость. Создадим оптимизатор с импульсом 0.9 и весовым затуханием 1e-4:

import torch

model = torch.nn.Linear(10, 5)
optimizer = torch.optim.SGD(
	model.parameters(),
	lr=0.01,
	momentum=0.9,
	weight_decay=1e-4
)
print(f"momentum: {optimizer.defaults['momentum']}")
print(f"weight_decay: {optimizer.defaults['weight_decay']}")

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

momentum: 0.9
weight_decay: 0.0001

Пример

Метод zero_grad обнуляет градиенты всех оптимизируемых параметров. Это важный шаг перед вычислением новых градиентов, чтобы избежать накопления градиентов от предыдущих итераций:

import torch

model = torch.nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

x = torch.tensor([[1.0, 2.0]])
y = torch.tensor([[1.0]])

criterion = torch.nn.MSELoss()

for epoch in range(2):
	optimizer.zero_grad()
	pred = model(x)
	loss = criterion(pred, y)
	loss.backward()
	optimizer.step()
	print(f"Epoch {epoch + 1}, Loss: {loss.item()}")

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

Epoch 1, Loss: 0.8130650511474609
Epoch 2, Loss: 0.7209780812263489

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

класс Adam,
который реализует адаптивный метод оценки моментов
класс RMSprop,
который использует среднеквадратичное распространение ошибки
класс Adagrad,
который адаптирует скорость обучения для каждого параметра
функцию clip_grad_norm_,
которая ограничивает норму градиентов для предотвращения взрыва