Пакет автоматической смешанной точности - torch.amp

Автоматическое переключение типов данных

class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None) [source]

Экземпляры autocast служат контекстными менеджерами или декораторами, которые позволяют определённым частям вашего скрипта выполняться в смешанной точности.

В этих областях операции выполняются в типе данных, специфичном для операции и выбранном по умолчанию для повышения производительности при сохранении точности. Подробности см. в Справочнике по операциям Autocast.

При входе в область с включённой автоматической смешанной точностью тензоры могут быть любого типа. Не следует вызывать half() или bfloat16() для вашей модели(ей) или входных данных при использовании автоматического переключения типов.

autocast должен охватывать только прямой(е) проход(ы) вашей сети, включая вычисление(я) потерь. Обратные проходы под управлением autocast не рекомендуются. Обратные операции выполняются в том же типе, что и соответствующие прямые операции, используемые autocast.

Пример для устройств CUDA:

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

См. примеры автоматической смешанной точности CUDA (вместе с масштабированием градиента) для использования в более сложных сценариях (например, штраф за градиент, несколько моделей/функций потерь, пользовательские функции autograd) на странице примеров.

autocast также может использоваться как декоратор, например, для метода forward вашей модели:

class AutocastModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
        ...

Числовые тензоры, полученные в области с включённой автоматической смешанной точностью, могут быть float16. После возвращения в область без автоматического переключения типов их использование с числовыми тензорами других типов может привести к ошибкам несоответствия типов. В таком случае преобразуйте полученные тензоры из области автокаста обратно в float32 (или в другой тип, если необходимо). Если тензор из области автокаста уже float32, преобразование является пустой операцией и не несёт дополнительной нагрузки. Пример для CUDA:

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    # torch.mm is on autocast's list of ops that should run in float16.
    # Inputs are float32, but the op runs in float16 and produces float16 output.
    # No manual casts are required.
    e_float16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_float16 = torch.mm(d_float32, e_float16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())

Пример обучения на процессоре:

# Creates model and optimizer in default precision
model = Net()
optimizer = optim.SGD(model.parameters(), ...)

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
            output = model(input)
            loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

Пример вывода на процессоре:

# Creates model in default precision
model = Net().eval()

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    for input in data:
        # Runs the forward pass with autocasting.
        output = model(input)

Пример вывода на процессоре с использованием Jit Trace:

class TestModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super(TestModel, self).__init__()
        self.fc1 = nn.Linear(input_size, num_classes)
    def forward(self, x):
        return self.fc1(x)

input_size = 2
num_classes = 2
model = TestModel(input_size, num_classes).eval()

# For now, we suggest to disable the Jit Autocast Pass,
# As the issue: https://github.com/pytorch/pytorch/issues/75956
torch._C._jit_set_autocast_mode(False)

with torch.cpu.amp.autocast(cache_enabled=False):
    model = torch.jit.trace(model, torch.randn(1, input_size))
model = torch.jit.freeze(model)
# Models Run
for _ in range(3):
    model(torch.randn(1, input_size))

Ошибки несоответствия типов в области с включённой автоматической смешанной точностью являются ошибкой; если вы наблюдаете это, пожалуйста, отправьте сообщение об ошибке.

autocast(enabled=False) подрегионы могут быть вложены в области с включённой автоматической смешанной точностью. Местное отключение автоматического переключения типов может быть полезно, например, если вы хотите принудительно выполнить подрегион в определённом dtype. Отключение автоматического переключения типов даёт вам явный контроль над типом выполнения. В подрегионе входные данные из окружающей области должны быть преобразованы в dtype перед использованием:

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(enabled=False):
        # Calls e_float16.float() to ensure float32 execution
        # (necessary because e_float16 was created in an autocasted region)
        f_float32 = torch.mm(c_float32, e_float16.float())

    # No manual casts are required when re-entering the autocast-enabled region.
    # torch.mm again runs in float16 and produces float16 output, regardless of input types.
    g_float16 = torch.mm(d_float32, f_float32)

Состояние автокаста является локальным для потока. Если вы хотите включить его в новом потоке, контекстный менеджер или декоратор должны вызываться в этом потоке. Это влияет на torch.nn.DataParallel и torch.nn.parallel.DistributedDataParallel при использовании более чем одного графического процессора на процесс (см. работу с несколькими графическими процессорами).

Параметры:

• device_type (str, обязательно) – Использовать ли устройство ‘cuda’ или ‘cpu’
• enabled (bool, необязательно) – Автоматический выбор типа данных должен быть включен в данной области. По умолчанию: True
• dtype (torch_dtype, необязательно) – Использовать ли torch.float16 или torch.bfloat16.
• cache_enabled (bool, необязательно) – Включить ли кеширование весов внутри autocast. По умолчанию: True

class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True) [source]

См. torch.autocast. torch.cuda.amp.autocast(args...) эквивалентно torch.autocast("cuda", args...)

torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None) [source]

Вспомогательный декоратор для методов forward пользовательских функций автоградиента (подклассы torch.autograd.Function). Дополнительные сведения см. на странице примера.

Параметры:

cast_inputs (torch.dtype или None, необязательно, по умолчанию=None) – Если не None, при выполнении forward в области с включённой автоматической смешанной точностью преобразование входных числовых тензоров CUDA в целевой тип (нечисловые тензоры не затрагиваются), затем выполнение forward с отключённой автоматической смешанной точностью. Если None, внутренние операции forward выполняются с текущим состоянием autocast.

Примечание

Если вызываемый декоратором forward метод вызывается вне области с включённой автоматической смешанной точностью, custom_fwd является пустой операцией, и cast_inputs не оказывает никакого влияния.

torch.cuda.amp.custom_bwd(bwd) [source]

Вспомогательный декоратор для методов обратного распространения пользовательских функций автоградиента (подклассы torch.autograd.Function). Гарантирует, что backward выполняется с тем же состоянием autocast, что и forward. Дополнительные сведения см. на странице примера.

class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True) [source]

См. torch.autocast. torch.cpu.amp.autocast(args...) эквивалентно torch.autocast("cpu", args...)

Масштабирование градиентов

Если прямой проход для конкретной операции имеет float16 входные данные, обратный проход для этой операции произведёт float16 градиенты. Значения градиентов с небольшой величиной могут быть не представимы в float16. Эти значения обнулятся («потеря значащих разрядов»), поэтому обновление соответствующих параметров будет утеряно.

Для предотвращения потери значащих разрядов, «масштабирование градиентов» умножает потерю(и) сети на масштабный множитель и вызывает обратный проход по масштабированной(ым) потерям. Градиенты, текущие назад по сети, затем масштабируются на тот же множитель. Другими словами, значения градиентов имеют большую величину, поэтому они не обнуляются.

Градиент каждого параметра (.grad атрибут) должен быть не масштабирован перед обновлением оптимизатором параметров, чтобы множитель масштабирования не мешал скорости обучения.

class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True) [source]

get_backoff_factor() [source]

Возвращает Python-число с коэффициентом отката масштабирования.

get_growth_factor() [source]

Возвращает Python-число с коэффициентом роста масштабирования.

get_growth_interval() [source]

Возвращает Python-целое число с интервалом роста.

get_scale() [source]

Возвращает Python-число с текущим масштабом или 1.0, если масштабирование отключено.

Предупреждение

get_scale() влечёт синхронизацию CPU-GPU.

is_enabled() [source]

Возвращает булево значение, указывающее, включен ли этот экземпляр.

load_state_dict(state_dict) [source]

Загружает состояние масштабирования. Если этот экземпляр отключён, load_state_dict() — это бездействие.

Параметры:

state_dict (dict) — состояние масштабирования. Должно быть объектом, возвращаемым из вызова state_dict().

scale(outputs) [source]

Умножает («масштабирует») тензор или список тензоров на множитель масштабирования.

Возвращает масштабированные результаты. Если этот экземпляр GradScaler не включен, результаты возвращаются без изменений.

Параметры:

outputs (Tensor или итерируемый список тензоров) — результаты, подлежащие масштабированию.

set_backoff_factor(new_factor) [source]

Параметры:

new_scale (float) — значение, используемое в качестве нового коэффициента отката масштабирования.

set_growth_factor(new_factor) [source]

Параметры:

new_scale (float) — значение, используемое в качестве нового коэффициента роста масштабирования.

set_growth_interval(new_interval) [source]

Параметры:

new_interval (int) — значение, используемое в качестве нового интервала роста.

state_dict() [source]

Возвращает состояние масштабирования в виде dict. Содержит пять элементов:

• "scale" — Python-число, содержащее текущий масштаб
• "growth_factor" — Python-число, содержащее текущий коэффициент роста
• "backoff_factor" — Python-число, содержащее текущий коэффициент отката
• "growth_interval" — Python-целое число, содержащее текущий интервал роста
• "_growth_tracker" — Python-целое число, содержащее количество недавних последовательных не пропущенных шагов.

Если этот экземпляр не включён, возвращает пустой словарь.

Примечание

Если вы хотите сохранить состояние масштабирования после определённой итерации, state_dict() следует вызвать после update().

step(optimizer, *args, **kwargs) [source]

step() выполняет следующие две операции:

1. Внутренне вызывает unscale_(optimizer) (если unscale_() не был явно вызван для optimizer ранее в итерации). В рамках unscale_() проверяются градиенты на наличие inf/NaN.
2. Если не обнаружены градиенты inf/NaN, вызывается optimizer.step() с использованием не масштабированных градиентов. В противном случае, optimizer.step() пропускается, чтобы избежать повреждения параметров.

*args и **kwargs передаются в optimizer.step().

Возвращает результат вызова optimizer.step(*args, **kwargs).

Параметры:

• optimizer (torch.optim.Optimizer) — оптимизатор, применяющий градиенты.
• args — любые аргументы.
• kwargs — любые ключевые аргументы.

Предупреждение

Использование замкнутых функций (closure) в настоящее время не поддерживается.

unscale_(optimizer) [source]

Делит («растягивает») тензоры градиентов оптимизатора на коэффициент масштабирования.

unscale_() является необязательным, обеспечивая случаи, когда вам нужно изменить или проверить градиенты между проходом(ами) обратного распространения и step(). Если unscale_() не вызывается явно, градиенты будут автоматически растягиваться во время step().

Простой пример, использующий unscale_() для включения обрезки нерастянутых градиентов:

...
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

Параметры:

optimizer (torch.optim.Optimizer) – Оптимизатор, которому принадлежат градиенты, подлежащие растяжению.

Примечание

unscale_() не требует синхронизации CPU-GPU.

Предупреждение

unscale_() следует вызывать только один раз на один оптимизатор на вызов step() и только после того, как все градиенты для назначенных параметров данного оптимизатора были накоплены. Вызов unscale_() дважды для данного оптимизатора между каждым вызовом step() приводит к ошибке RuntimeError.

Предупреждение

unscale_() может растянуть разреженные градиенты вне места, заменяя атрибут .grad.

update(new_scale=None) [source]

Обновляет коэффициент масштабирования.

Если какие-либо шаги оптимизатора были пропущены, масштаб умножается на backoff_factor для уменьшения. Если growth_interval непропущенных итераций произошли последовательно, масштаб умножается на growth_factor для увеличения.

Передача new_scale задаёт новое значение масштаба вручную. (new_scale не используется напрямую, он используется для заполнения внутреннего тензора масштаба GradScaler. Таким образом, если new_scale был тензором, последующие изменения этого тензора in-place не будут далее влиять на масштаб, используемый GradScaler внутри.)

Параметры:

new_scale (float or torch.cuda.FloatTensor, optional, default=None) – Новый коэффициент масштабирования.

Предупреждение

update() следует вызывать только в конце итерации, после того, как scaler.step(optimizer) был вызван для всех оптимизаторов, используемых в этой итерации.

Справочник по операциям Autocast