Метод scale

Метод scale класса GradScaler используется для масштабирования градиентов во время обучения с использованием автоматического смешанного обучения (AMP). Он принимает на вход один или несколько тензоров градиентов и возвращает их масштабированные версии. Масштаб автоматически подбирается для предотвращения андерфлоу чисел с плавающей точкой.

Синтаксис

scaler.scale(tensors)

Параметры

Метод принимает один или несколько тензоров градиентов в качестве позиционных аргументов. Это могут быть тензоры, списки тензоров или словари с тензорами.

loss1 = output1 - target1
loss2 = output2 - target2
scaler.scale(loss1).backward()
scaler.scale(loss2).backward()

Метод возвращает масштабированный тензор или структуру, соответствующую переданной, с тем же типом данных. Результат можно сразу использовать для вызова backward.

Пример применения с одним тензором

Рассмотрим базовое использование метода для масштабирования потерь при обучении модели:

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)
model = torch.nn.Linear(10, 1)
data = torch.randn(5, 10)
target = torch.randn(5, 1)
scaler = GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

with torch.cuda.amp.autocast():
	output = model(data)
	loss = torch.nn.functional.mse_loss(output, target)

scaler.scale(loss).backward()
print("Масштабирование выполнено успешно")

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

"Масштабирование выполнено успешно"

Пример с несколькими тензорами

Метод может масштабировать несколько тензоров потерь одновременно:

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)
model = torch.nn.Linear(10, 5)
data = torch.randn(3, 10)
target1 = torch.randn(3, 5)
target2 = torch.randn(3, 5)
scaler = GradScaler()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

with torch.cuda.amp.autocast():
	output = model(data)
	loss1 = torch.nn.functional.mse_loss(output, target1)
	loss2 = torch.nn.functional.mse_loss(output, target2)

loss = loss1 + loss2
scaler.scale(loss).backward()
print("Масштабирование потерь выполнено")

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

"Масштабирование потерь выполнено"

Пример в цикле обучения

Покажем полный цикл обучения с использованием метода в связке с другими методами класса:

import torch
from torch.cuda.amp import GradScaler, autocast

torch.manual_seed(0)
model = torch.nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler()
data = torch.randn(100, 10)
target = torch.randn(100, 1)

for epoch in range(3):
	for i in range(0, 100, 10):
		batch_data = data[i:i+10]
		batch_target = target[i:i+10]
		
		optimizer.zero_grad()
		
		with autocast():
			output = model(batch_data)
			loss = torch.nn.functional.mse_loss(output, batch_target)
		
		scaler.scale(loss).backward()
		scaler.step(optimizer)
		scaler.update()
	
	print(f"Эпоха {epoch + 1} завершена")

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

"Эпоха 1 завершена"
"Эпоха 2 завершена"
"Эпоха 3 завершена"

Пример с обработкой ошибок

Метод корректно обрабатывает ситуации, когда градиенты содержат значения, близкие к нулю:

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)
model = torch.nn.Linear(5, 1)
scaler = GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
data = torch.randn(2, 5)
target = torch.zeros(2, 1)

with torch.cuda.amp.autocast():
	output = model(data)
	loss = torch.nn.functional.mse_loss(output, target)

scaler.scale(loss).backward()
print("Масштабирование выполнено без ошибок")

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

"Масштабирование выполнено без ошибок"

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

класс GradScaler,
который управляет масштабированием градиентов в AMP
метод step,
который выполняет шаг оптимизатора с учетом масштаба
метод update,
который обновляет коэффициент масштабирования
метод unscale_,
который отменяет масштабирование градиентов