Метод unscale_

Метод unscale_ класса GradScaler отменяет масштабирование градиентов для переданных тензоров. Этот метод вызывается перед обновлением весов оптимизатором, чтобы градиенты были в правильном масштабе. Метод работает только с тензорами, которые были созданы с использованием автоматического смешанного точности (AMP) и масштабированы методом scale.

Метод принимает на вход оптимизатор, градиенты которого необходимо отмасштабировать. Важно вызывать этот метод перед вызовом step оптимизатора и только один раз для каждого оптимизатора за итерацию.

Синтаксис

grad_scaler.unscale_(optimizer)

Пример

Базовый пример использования метода unscale_ при обучении модели с автоматическим смешанным точностью:

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler

# Создание модели и оптимизатора
model = nn.Linear(10, 10).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler()

# Фиктивные данные
data = torch.randn(8, 10).cuda()
target = torch.randn(8, 10).cuda()

# Прямой проход с автоматическим смешанным точность
with torch.cuda.amp.autocast():
	output = model(data)
	loss = nn.MSELoss()(output, target)

# Обратный проход и масштабирование градиентов
scaler.scale(loss).backward()

# Отмена масштабирования градиентов
scaler.unscale_(optimizer)

# Обновление весов модели
scaler.step(optimizer)
scaler.update()

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

"Модель успешно обучена с использованием AMP"

Пример

Пример использования метода unscale_ перед применением клиппинга градиентов:

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler
import torch.nn.utils as utils

# Создание модели и оптимизатора
model = nn.Linear(10, 10).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler()

# Фиктивные данные
data = torch.randn(8, 10).cuda()
target = torch.randn(8, 10).cuda()

# Прямой и обратный проходы
with torch.cuda.amp.autocast():
	output = model(data)
	loss = nn.MSELoss()(output, target)

scaler.scale(loss).backward()

# Отмена масштабирования перед клиппингом
scaler.unscale_(optimizer)

# Применение клиппинга градиентов
utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Обновление весов модели
scaler.step(optimizer)
scaler.update()

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

"Градиенты отмасштабированы и применен клиппинг"

Пример

Пример использования метода unscale_ с несколькими оптимизаторами:

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler

# Создание модели с двумя оптимизаторами
model = nn.Sequential(
	nn.Linear(10, 10),
	nn.ReLU(),
	nn.Linear(10, 10)
).cuda()

# Разделение параметров на две группы
params1 = list(model[0].parameters())
params2 = list(model[2].parameters())

optimizer1 = torch.optim.SGD(params1, lr=0.01)
optimizer2 = torch.optim.SGD(params2, lr=0.02)
scaler = GradScaler()

# Фиктивные данные
data = torch.randn(8, 10).cuda()
target = torch.randn(8, 10).cuda()

# Прямой и обратный проходы
with torch.cuda.amp.autocast():
	output = model(data)
	loss = nn.MSELoss()(output, target)

scaler.scale(loss).backward()

# Отмена масштабирования для каждого оптимизатора
scaler.unscale_(optimizer1)
scaler.unscale_(optimizer2)

# Обновление весов моделей
scaler.step(optimizer1)
scaler.step(optimizer2)
scaler.update()

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

"Модель обновлена с использованием двух оптимизаторов"

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

класс GradScaler,
который управляет масштабированием градиентов для обучения
метод scale,
который применяет масштабирование к потерям перед обратным проходом
метод step,
который выполняет шаг оптимизации с учетом масштабирования
метод update,
который обновляет масштабирующий множитель в конце итерации