Класс GradScaler
Класс GradScaler из модуля torch.cuda.amp
предназначен для автоматического масштабирования градиентов
при обучении нейронных сетей с использованием
смешанной точности (mixed precision).
Основная задача класса - умножать потери на масштабирующий
коэффициент перед обратным распространением, чтобы
предотвратить обнуление градиентов при использовании
типа float16. После обновления весов масштабирующий
коэффициент автоматически корректируется, что позволяет
поддерживать стабильность обучения.
Синтаксис
torch.cuda.amp.GradScaler(
init_scale=2.0**16,
growth_factor=2.0,
backoff_factor=0.5,
growth_interval=2000,
enabled=True
)
Основные параметры конструктора:
-
init_scale- начальное значение масштабирующего коэффициента (по умолчанию 65536) -
growth_factor- множитель для увеличения масштаба при успешном обучении (по умолчанию 2.0) -
backoff_factor- множитель для уменьшения масштаба при обнаружении переполнения (по умолчанию 0.5) -
growth_interval- количество шагов между попытками увеличения масштаба (по умолчанию 2000) -
enabled- флаг включения/отключения масштабирования (по умолчаниюTrue)
Пример
Давайте создадим экземпляр класса GradScaler
с настройками по умолчанию:
import torch
scaler = torch.cuda.amp.GradScaler()
print(scaler)
Результат выполнения кода:
<torch.cuda.amp.grad_scaler.GradScaler object at 0x7f8a2c3d4e50>
Пример
Создадим масштабировщик с пользовательскими параметрами:
import torch
scaler = torch.cuda.amp.GradScaler(
init_scale=2.0**10,
growth_factor=1.5,
backoff_factor=0.25,
growth_interval=100
)
print(f"Initial scale: {scaler.get_scale()}")
Результат выполнения кода:
"Initial scale: 1024.0"
Пример
Базовый цикл обучения с использованием масштабировщика для смешанной точности:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = torch.cuda.amp.GradScaler()
x = torch.randn(32, 10)
y = torch.randn(32, 1)
for epoch in range(2):
optimizer.zero_grad()
with torch.cuda.amp.autocast():
pred = model(x)
loss = nn.MSELoss()(pred, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Epoch {epoch+1}: scale={scaler.get_scale()}")
Результат выполнения кода:
"Epoch 1: scale=65536.0"
"Epoch 2: scale=65536.0"
Пример
Пример использования масштабировщика с отключённым масштабированием (полезно для отладки):
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(5, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler(enabled=False)
x = torch.randn(16, 5)
y = torch.randint(0, 2, (16,))
optimizer.zero_grad()
with torch.cuda.amp.autocast(enabled=False):
pred = model(x)
loss = nn.CrossEntropyLoss()(pred, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Scaler enabled: {scaler.is_enabled()}")
Результат выполнения кода:
"Scaler enabled: False"
