Ограничение нормы градиента в PyTorch

Если норма градиента по всем параметрам слишком велика, шаг обновления может разрушить стабильность модели. После обратного прохода и до вызова шага у оптимизатора градиенты можно сжать через функцию clip_grad_norm_ из torch.nn.utils.

Функция возвращает норму до обрезки. Зафиксируем зерно, считаем ошибку на случайном батче и выведем это число:

import torch
import torch.nn as nn

torch.manual_seed(42)

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

x = torch.randn(4, 3)
y = torch.randn(4, 1)

optimizer.zero_grad()
loss = nn.MSELoss()(model(x), y)
loss.backward()

total_norm = torch.nn.utils.clip_grad_norm_(
	model.parameters(),
	max_norm=1.0,
)
print(total_norm) # выведет tensor(2.8757)

optimizer.step()

Соберите линейный слой 4 на 1 и спуск со скоростью 0.01. После обратного прохода по случайным данным ограничьте норму градиента порогом 0.5 и выведите число, которое вернула функция обрезки.

При зерне 0 обучите однослойную модель на 8 примерах, обрежите градиент с порогом 2.0, сделайте шаг оптимизатора и выведите только норму до обрезки.

Возьмите два параметра линейного слоя 2 на 2, заполните их единицами, задайте градиенты вручную как 3 и 4. Ограничьте суммарную норму 1 и выведите норму, которую вернула обрезка.