Масштаб ошибки в PyTorch

Класс GradScaler из torch.amp масштабирует ошибку перед обратным проходом и возвращает шаг оптимизатора к исходному масштабу. Так защищают обратный проход при половинной точности на видеокарте.

Без видеокарты масштаб обычно не включают: объект создают для процессора и проверяют, активен ли он:

import torch
import torch.nn as nn

scaler = torch.amp.GradScaler("cpu", enabled=False)
print(scaler.is_enabled()) # выведет False

Когда масштаб выключен, цепочка scale, backward, step и update всё равно может отработать на маленьком примере:

import torch
import torch.nn as nn

weight = torch.tensor(2.0, requires_grad=True)
scaler = torch.amp.GradScaler("cpu", enabled=False)
loss = (weight * 3).sum()
scaler.scale(loss).backward()
optimizer = torch.optim.SGD([weight], lr=0.1)
scaler.step(optimizer)
scaler.update()
print(weight.item()) # выведет 1.7000000476837158

Создайте масштаб для процессора без включения и выведите, активен ли он.

Для числа 1.0 с записью градиента прогоните один шаг с масштабом на процессоре и выведите значение параметра после обновления.

Сравните два масштаба на процессоре: один создайте без включения, второй с явным включением для процессора. Выведите два признака активности через пробел.