Класс SparseAdam
Класс SparseAdam является реализацией алгоритма Adam, адаптированной для работы с разреженными тензорами и градиентами. Этот оптимизатор эффективно обрабатывает случаи, когда большинство параметров модели имеют нулевые градиенты, что характерно для рекомендательных систем, моделей обработки естественного языка с большими словарями и других задач с разреженными признаками. В отличие от стандартного Adam, SparseAdam выполняет обновления только для тех параметров, для которых градиент был ненулевым.
Класс принимает следующие основные параметры: params (итерируемый объект параметров или словарей с параметрами), lr (скорость обучения, по умолчанию 1e-3), betas (кортеж коэффициентов затухания для моментов, по умолчанию (0.9, 0.999)), eps (член стабилизации, по умолчанию 1e-8) и weight_decay (коэффициент регуляризации L2, по умолчанию 0).
Синтаксис
torch.optim.SparseAdam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0)
Пример
Рассмотрим базовый пример использования SparseAdam для оптимизации разреженных параметров:
import torch
import torch.nn as nn
from torch.optim import SparseAdam
torch.manual_seed(0)
model = nn.Embedding(10, 5, sparse=True)
optimizer = SparseAdam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
indices = torch.tensor([1, 3, 5])
x = torch.tensor([[1], [3], [5]])
y = torch.tensor([0, 2, 1])
optimizer.zero_grad()
out = model(x).squeeze(1)
loss = loss_fn(out, y)
loss.backward()
optimizer.step()
print(model.weight[1])
Результат выполнения кода:
tensor([ 0.0881, 0.0712, 0.0042, -0.0990, 0.0658],
grad_fn=<SelectBackward0>)
Пример
Продемонстрируем обновление только разреженных параметров при нескольких шагах оптимизации:
import torch
import torch.nn as nn
from torch.optim import SparseAdam
torch.manual_seed(0)
embedding = nn.Embedding(10, 3, sparse=True)
optimizer = SparseAdam(embedding.parameters(), lr=0.01)
print("Before:", embedding.weight[2])
indices = torch.tensor([2])
x = torch.tensor([[2]])
loss = embedding(x).sum()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("After:", embedding.weight[2])
Результат выполнения кода:
Before: tensor([-1.1248, -0.3667, -0.5292], grad_fn=<SelectBackward0>)
After: tensor([-1.1335, -0.3740, -0.5381], grad_fn=<SelectBackward0>)
Пример
Сравним поведение SparseAdam с Adam на разреженных данных. При одинаковых начальных условиях обновления будут выполняться только для затронутых параметров:
import torch
import torch.nn as nn
from torch.optim import SparseAdam, Adam
torch.manual_seed(0)
embed_sparse = nn.Embedding(10, 4, sparse=True)
embed_dense = nn.Embedding(10, 4, sparse=False)
embed_dense.weight.data = embed_sparse.weight.data.clone()
opt_sparse = SparseAdam(embed_sparse.parameters(), lr=0.01)
opt_dense = Adam(embed_dense.parameters(), lr=0.01)
idx = torch.tensor([3])
x = torch.tensor([[3]])
opt_sparse.zero_grad()
loss_sparse = embed_sparse(x).sum()
loss_sparse.backward()
opt_sparse.step()
opt_dense.zero_grad()
loss_dense = embed_dense(x).sum()
loss_dense.backward()
opt_dense.step()
print("SparseAdam update:", embed_sparse.weight[3])
print("Adam update:", embed_dense.weight[3])
Результат выполнения кода:
SparseAdam update: tensor([-0.6155, -1.1746, -0.7284, -0.0503],
grad_fn=<SelectBackward0>)
Adam update: tensor([-0.6160, -1.1751, -0.7288, -0.0507],
grad_fn=<SelectBackward0>)
