Метод update

Метод update класса GradScaler обновляет состояние масштабирования после завершения шага оптимизации. Он вызывается после методов scale и step и принимает на вход оптимизатор, который был использован для обновления весов.

Синтаксис

scaler.update(optimizer)

Метод update анализирует градиенты, полученные в результате обратного распространения, и при необходимости корректирует значение масштабирующего множителя. Если градиенты содержат значения NaN или Inf, масштабирующий множитель уменьшается. В противном случае он может быть увеличен до максимально допустимого значения.

Пример

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

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler()

data = torch.randn(5, 10)
target = torch.randn(5, 1)

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

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update(optimizer)

print("Обновление масштабирования выполнено")

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

"Обновление масштабирования выполнено"

Пример

В этом примере показан полный цикл обучения с использованием метода update внутри цикла по эпохам:

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scaler = GradScaler()

data = torch.randn(20, 10)
target = torch.randn(20, 1)

for epoch in range(3):
	optimizer.zero_grad()
	
	with torch.cuda.amp.autocast():
		output = model(data)
		loss = torch.nn.functional.mse_loss(output, target)
	
	scaler.scale(loss).backward()
	scaler.step(optimizer)
	scaler.update(optimizer)
	
	print(f"Epoch {epoch + 1}: loss = {loss.item():.4f}")

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

Epoch 1: loss = 0.9660
Epoch 2: loss = 0.9395
Epoch 3: loss = 0.9148

Пример

Метод update также можно использовать без передачи оптимизатора, если необходимо обновить состояние масштабирования без привязки к конкретному оптимизатору:

import torch
from torch.cuda.amp import GradScaler

torch.manual_seed(0)

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

data = torch.randn(4, 5)
target = torch.randn(4, 3)

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

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

print("Обновление выполнено без оптимизатора")

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

"Обновление выполнено без оптимизатора"

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

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