Функция clip_grad_norm_

Функция clip_grad_norm_ обрезает градиенты параметров модели по норме. Первым параметром функция принимает итератор параметров модели. Вторым параметром передается максимальная норма. Третьим параметром можно задать тип нормы. Функция изменяет градиенты на месте и возвращает общую норму градиентов.

Синтаксис

torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=2.0)

Пример

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

import torch
import torch.nn as nn

torch.manual_seed(0)
model = 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 = nn.MSELoss()
loss = criterion(model(x), y)
loss.backward()

total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
print(total_norm)

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

tensor(0.3388)

Пример

Давайте используем обрезание по норме в процессе обучения:

import torch
import torch.nn as nn

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

for epoch in range(5):
	x = torch.randn(4, 2)
	y = torch.randn(4, 1)
	
	optimizer.zero_grad()
	loss = nn.MSELoss()(model(x), y)
	loss.backward()
	
	torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
	optimizer.step()
	
	print(f'Epoch {epoch + 1}, loss: {loss.item():.4f}')

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

Epoch 1, loss: 2.5166
Epoch 2, loss: 1.3200
Epoch 3, loss: 2.1503
Epoch 4, loss: 0.4488
Epoch 5, loss: 1.4112

Пример

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

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(
	nn.Linear(2, 3),
	nn.ReLU(),
	nn.Linear(3, 1)
)

x = torch.randn(2, 2)
y = torch.randn(2, 1)

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

total_norm = torch.nn.utils.clip_grad_norm_(
	model.parameters(), 
	max_norm=1.0, 
	norm_type=float('inf')
)
print(total_norm)

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

tensor(1.0000)

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

функцию clip_grad_value_,
которая обрезает градиенты по абсолютному значению
функцию Adam,
которая является оптимизатором с адаптивным моментом
функцию SGD,
которая является базовым оптимизатором стохастического градиентного спуска
функцию CrossEntropyLoss,
которая вычисляет кросс-энтропийную функцию потерь