Атрибут state

Атрибут state класса Optimizer представляет собой словарь, который хранит состояние оптимизатора для каждого параметра модели. Ключами этого словаря являются тензоры параметров, а значениями - словари с данными состояния, такими как момент (momentum), квадрат градиента и другие величины, необходимые для работы алгоритмов оптимизации. Этот атрибут автоматически обновляется при вызове метода step.

Синтаксис

optimizer.state

Пример

Давайте создадим простую модель и оптимизатор, чтобы посмотреть на содержимое атрибута state:

import torch
import torch.nn as nn

# Create a simple model
model = nn.Linear(10, 5)

# Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Create random input and target
x = torch.randn(1, 10)
y = torch.randn(1, 5)

# Forward pass
output = model(x)
loss = torch.nn.functional.mse_loss(output, y)

# Backward pass
loss.backward()

# Perform one optimization step
optimizer.step()

# Print state
print(optimizer.state)

Результат выполнения кода (в зависимости от случайных чисел может отличаться):

{
	Parameter containing:
	tensor([[-0.1526,  0.2961,  0.2458,  0.2574, -0.1409, -0.0104, -0.0170,  0.2114,
			-0.0127, -0.1758],
			[ 0.0828, -0.0095,  0.1054, -0.0018, -0.1010, -0.1793,  0.0688, -0.0401,
			-0.2299, -0.0565],
			[-0.0349,  0.2262,  0.0105, -0.1918, -0.0344, -0.1532, -0.1268,  0.1559,
			0.0118, -0.1772],
			[-0.0763,  0.2706,  0.1610, -0.0238,  0.0532, -0.1713,  0.1577, -0.0885,
			0.1941, -0.1104],
			[ 0.0232, -0.2817, -0.1190,  0.0105,  0.0991, -0.2790, -0.2623,  0.2133,
			-0.0543,  0.2633]], requires_grad=True): {
				'momentum_buffer': tensor([[...]])
			}
}

Пример

Давайте рассмотрим, как атрибут state изменяется после вызова метода zero_grad и step:

import torch
import torch.nn as nn

# Set random seed for reproducibility
torch.manual_seed(0)

# Create a simple model
model = nn.Linear(5, 3)

# Create Adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Generate random data
x = torch.randn(2, 5)
y = torch.randn(2, 3)

# Initial state before any step
print("Initial state:")
print(optimizer.state)

# Forward and backward passes
output = model(x)
loss = torch.nn.functional.mse_loss(output, y)
loss.backward()

# Perform optimization step
optimizer.step()

# State after one step
print("\nState after step:")
print(optimizer.state)

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

Initial state:
{}

State after step:
{
	Parameter containing:
	tensor([[...]], requires_grad=True): {
		'exp_avg': tensor([[...]]),
		'exp_avg_sq': tensor([[...]]),
		'step': tensor(1)
	}
}

Пример

Давайте рассмотрим, как можно получить доступ к состоянию для конкретного параметра:

import torch
import torch.nn as nn

# Create model and optimizer
model = nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Generate random data
x = torch.randn(3, 10)
y = torch.randn(3, 5)

# Training step
output = model(x)
loss = torch.nn.functional.mse_loss(output, y)
loss.backward()
optimizer.step()

# Get state for a specific parameter
for param in model.parameters():
	if param in optimizer.state:
		param_state = optimizer.state[param]
		print(f"Parameter shape: {param.shape}")
		print(f"State keys: {param_state.keys()}")
		break

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

Parameter shape: torch.Size([5, 10])
State keys: dict_keys(['momentum_buffer'])

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

класс Optimizer,
который является базовым классом для всех оптимизаторов в PyTorch
метод step,
который выполняет один шаг оптимизации и обновляет состояние
атрибут param_groups,
который содержит группы параметров с их настройками
метод state_dict,
который возвращает словарь с полным состоянием оптимизатора