Метод step

Метод step класса Optimizer выполняет обновление параметров модели, используя вычисленные градиенты. Этот метод вызывается после выполнения обратного распространения ошибки и является ключевым этапом в процессе обучения нейронных сетей. При вызове step оптимизатор применяет выбранный алгоритм оптимизации (SGD, Adam, RMSprop и другие) для обновления весов модели.

Синтаксис

optimizer.step()

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

optimizer.step(closure)

Пример

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

import torch
import torch.nn as nn

torch.manual_seed(0)

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

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

for epoch in range(100):
	pred = model(x)
	loss = nn.MSELoss()(pred, y)
	
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

print(model.weight.item(), model.bias.item())

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

2.0001156321258545 0.000600355790939182

Пример

Использование оптимизатора Adam с методом step для обучения модели на более сложных данных:

import torch
import torch.nn as nn

torch.manual_seed(0)

model = nn.Sequential(
	nn.Linear(2, 4),
	nn.ReLU(),
	nn.Linear(4, 1)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

x = torch.tensor([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]])
y = torch.tensor([[3.0], [5.0], [7.0]])

for epoch in range(200):
	pred = model(x)
	loss = nn.MSELoss()(pred, y)
	
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

print(loss.item())

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

0.00014271571962260455

Пример

Использование параметра closure для алгоритма LBFGS, который требует повторного вычисления потерь:

import torch
import torch.nn as nn

torch.manual_seed(0)

model = nn.Linear(1, 1)
optimizer = torch.optim.LBFGS(model.parameters(), lr=1.0)

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

for epoch in range(100):
	def closure():
		optimizer.zero_grad()
		pred = model(x)
		loss = nn.MSELoss()(pred, y)
		loss.backward()
		return loss
	
	optimizer.step(closure)

print(model.weight.item(), model.bias.item())

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

2.0000100135803223 0.00019091507430928946

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

класс Optimizer,
который является базовым классом для всех оптимизаторов
метод zero_grad,
который обнуляет градиенты перед обратным распространением
метод state_dict,
который возвращает состояние оптимизатора для сохранения
метод load_state_dict,
который загружает состояние оптимизатора из сохранённого словаря