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

torch.amp предоставляет удобные методы для смешанной точности, где некоторые операции используют тип данных torch.float32 (float) и другие операции используют тип данных с плавающей точкой меньшей точности (lower_precision_fp): torch.float16 (half) или torch.bfloat16. Некоторые операции, такие как линейные слои и свертки, намного быстрее в lower_precision_fp. Другие операции, такие как сокращения, часто требуют динамического диапазона float32. Смешанная точность пытается сопоставить каждую операцию с соответствующим типом данных.

Обычно «обучение со смешанной точностью» с типом данных torch.float16 использует torch.autocast и torch.cuda.amp.GradScaler вместе, как показано в примерах CUDA автоматической смешанной точности и рецепте CUDA автоматической смешанной точности. Однако torch.autocast и torch.cuda.amp.GradScaler являются модульными и могут использоваться по отдельности, если это необходимо. Как показано в разделе примера для CPU в torch.autocast, «обучение/вывод со смешанной точностью» на CPU с типом данных torch.bfloat16 использует только torch.autocast.

Для CUDA и CPU также предоставляются отдельные API:

• torch.autocast("cuda", args...) эквивалентно torch.cuda.amp.autocast(args...).
• torch.autocast("cpu", args...) эквивалентно torch.cpu.amp.autocast(args...). Для CPU в настоящее время поддерживается только тип данных с плавающей точкой меньшей точности torch.bfloat16.

torch.autocast и torch.cpu.amp.autocast являются новыми в версии 1.10.

• Автоматическое преобразование типа
• Масштабирование градиента

• Справочник по операциям автоматического преобразования типа

  • Пригодность операций

  • Поведение операций CUDA

    • Операции CUDA, которые могут автоматически преобразовывать тип к float16
    • Операции CUDA, которые могут автоматически преобразовывать тип к float32
    • Операции CUDA, которые повышают тип до самого широкого типа входных данных
    • Предпочитать binary_cross_entropy_with_logits над binary_cross_entropy

  • Поведение операций CPU

    • Операции CPU, которые могут автоматически преобразовывать тип к bfloat16
    • Операции CPU, которые могут автоматически преобразовывать тип к float32
    • Операции CPU, которые повышают тип до самого широкого типа входных данных

Автоматическое преобразование типа

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

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

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

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

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 torch.autocast(device_type="cuda"):
        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):
    ...
    @torch.autocast(device_type="cuda")
    def forward(self, input):
        ...

Тензоры с плавающей точкой, созданные в области, поддерживаемой autocast, могут быть float16. После возвращения в область, не поддерживаемую autocast, использование их с тензорами с плавающей точкой разных типов может привести к ошибкам несоответствия типов. В таком случае преобразуйте тензор(ы), созданный в области autocast, обратно к float32 (или другому типу, если необходимо). Если тензор из области autocast уже 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 torch.autocast(device_type="cuda"):
    # 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())

Пример обучения на CPU:

# 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()

Пример вывода на CPU:

# 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)

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

class TestModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super().__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, являются ошибкой; если вы это наблюдаете, пожалуйста, откройте вопрос.

autocast(enabled=False) под-области могут быть вложены в области, поддерживаемые autocast. Местное отключение autocast может быть полезно, например, если вы хотите заставить под-область выполняться в конкретном dtype. Отключение autocast дает вам явный контроль над типом выполнения. В под-области входные данные из окружающей области должны быть приведены к типу 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 torch.autocast(device_type="cuda"):
    e_float16 = torch.mm(a_float32, b_float32)
    with torch.autocast(device_type="cuda", 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)

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

Parameters

• device_type (str, required) – Тип устройства для использования. Возможные значения: ‘cuda’, ‘cpu’, ‘xpu’ и ‘hpu’. Тип совпадает с атрибутом type torch.device. Таким образом, вы можете получить тип устройства тензора, используя Tensor.device.type.
• enabled (bool, optional) – Включить ли autocasting в области. По умолчанию: True
• dtype (torch_dtype, optional) – Используть ли torch.float16 или torch.bfloat16.
• cache_enabled (bool, optional) – Включить ли кеш весов внутри 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 выполняются с текущим состоянием автоматической конвертации.

Примечание

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

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

Функция-декоратор для методов обратного распространения пользовательских функций автоградиента (подклассы torch.autograd.Function). Обеспечивает, что backward выполняется с тем же состоянием автоматической конвертации, что и 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 float, содержащее коэффициент уменьшения масштаба.

get_growth_factor() [source]

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

get_growth_interval() [source]

Возвращает значение Python int, содержащее интервал увеличения.

get_scale() [source]

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

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

get_scale() требует синхронизации процессора и графического процессора.

is_enabled() [source]

Возвращает bool, указывающий, включена ли эта инстанция.

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 float, содержащее текущий масштаб
• "growth_factor" - значение Python float, содержащее текущий коэффициент увеличения
• "backoff_factor" - значение Python float, содержащее текущий коэффициент уменьшения масштаба
• "growth_interval" - значение Python int, содержащее текущий интервал увеличения
• "_growth_tracker" - значение Python int, содержащее количество последних последовательных шагов, которые не были пропущены.

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

Примечание

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

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

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

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

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

Возвращает значение, возвращаемое optimizer.step(*args, **kwargs).

Параметры

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

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

Использование замыканий в настоящее время не поддерживается.

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 был тензором, последующие изменения этого тензора не повлияют на используемый GradScaler масштаб.)

Параметры

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

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

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

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

По соображениям производительности, мы не проверяем значение коэффициента масштабирования, чтобы избежать синхронизаций, поэтому коэффициент масштабирования не гарантируется, что будет выше 1. Если масштаб опустится ниже 1 и/или вы видите NaN в ваших градиентах или потере, скорее всего, есть проблема. Например, предварительно обученные модели bf16 часто несовместимы с AMP/fp16 из-за разных динамических диапазонов.

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

Пригодность операций

Операции, которые выполняются в float64 или в нечисловых типах данных, не пригодны и будут выполняться в этих типах, независимо от того, включена ли функция автокастинга.

Пригодны только операции вне места и методы тензора. Внутри операции и вызовы с явным out=... тензором разрешены в областях, активированных функцией автокастинга, но не будут подвергаться автокастингу. Например, в области с включённым автокастингом a.addmm(b, c) может подвергнуться автокастингу, но a.addmm_(b, c) и a.addmm(b, c, out=d) не могут. Для лучшей производительности и стабильности предпочтительнее использовать операции вне места в активированных областях автокастинга.

Операции, вызванные с явным аргументом dtype=... не пригодны и будут производить вывод, соответствующий аргументу dtype.

Особенности операций CUDA

Следующие списки описывают поведение пригодных операций в активированных областях автокастинга. Эти операции всегда проходят автокастинг, независимо от того, вызваны ли они как часть torch.nn.Module, как функция или как метод torch.Tensor. Если функции экспонированы в нескольких пространствах имён, они подвергаются автокастингу независимо от пространства имён.

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

Если операция не указана, мы предполагаем, что она численно устойчива в float16. Если вы считаете, что неуказанная операция численно неустойчива в float16, пожалуйста, создайте вопрос.

Операции CUDA, которые могут быть переведены в float16

__matmul__, addbmm, addmm, addmv, addr, baddbmm, bmm, chain_matmul, multi_dot, conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, GRUCell, linear, LSTMCell, matmul, mm, mv, prelu, RNNCell

Операции CUDA, которые могут быть переведены в float32

__pow__, __rdiv__, __rpow__, __rtruediv__, acos, asin, binary_cross_entropy_with_logits, cosh, cosine_embedding_loss, cdist, cosine_similarity, cross_entropy, cumprod, cumsum, dist, erfinv, exp, expm1, group_norm, hinge_embedding_loss, kl_div, l1_loss, layer_norm, log, log_softmax, log10, log1p, log2, margin_ranking_loss, mse_loss, multilabel_margin_loss, multi_margin_loss, nll_loss, norm, normalize, pdist, poisson_nll_loss, pow, prod, reciprocal, rsqrt, sinh, smooth_l1_loss, soft_margin_loss, softmax, softmin, softplus, sum, renorm, tan, triplet_margin_loss

Операции CUDA, которые повышают до самого широкого типа входных данных

Эти операции не требуют конкретного типа данных для стабильности, но принимают несколько входных данных и требуют, чтобы типы данных входных данных совпадали. Если все входные данные являются float16, операция выполняется в float16. Если хотя бы один из входов является float32, автокастинг преобразует все входы в float32 и выполняет операцию в float32.

addcdiv, addcmul, atan2, bilinear, cross, dot, grid_sample, index_put, scatter_add, tensordot

Некоторые операции, не перечисленные здесь (например, бинарные операции, такие как add) изначально повышают входные данные без вмешательства автокастинга. Если входные данные являются смесью float16 и float32, эти операции выполняются в float32 и производят float32 выход, независимо от того, включён ли автокастинг.

Предпочитайте binary_cross_entropy_with_logits binary_cross_entropy

Обратные проходы torch.nn.functional.binary_cross_entropy() (и torch.nn.BCELoss, которая его оборачивает) могут создавать градиенты, которые не могут быть представлены в float16. В активированных областях автокастинга входное значение может быть float16, что означает, что обратный градиент должен быть представлен в float16 (автокастинг float16 входных данных в float32 не помогает, потому что это преобразование должно быть обращено вспять при обратном проходе). Поэтому, binary_cross_entropy и BCELoss выдают ошибку в активированных областях автокастинга.

Многие модели используют слой сигмоиды прямо перед слоем бинарной кросс-энтропии. В этом случае объедините два слоя, используя torch.nn.functional.binary_cross_entropy_with_logits() или torch.nn.BCEWithLogitsLoss. binary_cross_entropy_with_logits и BCEWithLogits можно безопасно использовать с autocast.

Поведение CPU-операций с автокастом

В следующих списках описывается поведение допустимых операций в областях с включенным автокастом. Эти операции всегда проходят через автокастинг, независимо от того, вызваны ли они как часть torch.nn.Module, как функция или как метод torch.Tensor. Если функции доступны в нескольких пространствах имен, они проходят через автокастинг независимо от пространства имен.

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

Если операция не указана, мы предполагаем, что она числово устойчива в bfloat16. Если вы считаете, что неуказанная операция числово неустойчива в bfloat16, пожалуйста, откройте вопрос.

CPU-операции, которые можно автокастировать в bfloat16

conv1d, conv2d, conv3d, bmm, mm, baddbmm, addmm, addbmm, linear, matmul, _convolution

CPU-операции, которые можно автокастировать в float32

conv_transpose1d, conv_transpose2d, conv_transpose3d, avg_pool3d, binary_cross_entropy, grid_sampler, grid_sampler_2d, _grid_sampler_2d_cpu_fallback, grid_sampler_3d, polar, prod, quantile, nanquantile, stft, cdist, trace, view_as_complex, cholesky, cholesky_inverse, cholesky_solve, inverse, lu_solve, orgqr, inverse, ormqr, pinverse, max_pool3d, max_unpool2d, max_unpool3d, adaptive_avg_pool3d, reflection_pad1d, reflection_pad2d, replication_pad1d, replication_pad2d, replication_pad3d, mse_loss, ctc_loss, kl_div, multilabel_margin_loss, fft_fft, fft_ifft, fft_fft2, fft_ifft2, fft_fftn, fft_ifftn, fft_rfft, fft_irfft, fft_rfft2, fft_irfft2, fft_rfftn, fft_irfftn, fft_hfft, fft_ihfft, linalg_matrix_norm, linalg_cond, linalg_matrix_rank, linalg_solve, linalg_cholesky, linalg_svdvals, linalg_eigvals, linalg_eigvalsh, linalg_inv, linalg_householder_product, linalg_tensorinv, linalg_tensorsolve, fake_quantize_per_tensor_affine, eig, geqrf, lstsq, _lu_with_info, qr, solve, svd, symeig, triangular_solve, fractional_max_pool2d, fractional_max_pool3d, adaptive_max_pool3d, multilabel_margin_loss_forward, linalg_qr, linalg_cholesky_ex, linalg_svd, linalg_eig, linalg_eigh, linalg_lstsq, linalg_inv_ex

CPU-операции, которые повышают до наибольшего типа входных данных

Эти операции не требуют конкретного типа данных для устойчивости, но принимают несколько входных данных и требуют, чтобы типы данных входных данных совпадали. Если все входные данные — bfloat16, операция выполняется в bfloat16. Если любой из входных данных — float32, автокаст приводит все входные данные к float32 и выполняет операцию в float32.

cat, stack, index_copy

Некоторые операции, не перечисленные здесь (например, бинарные операции, такие как add), по умолчанию повышают входные данные без вмешательства автокаста. Если входные данные являются смесью bfloat16 и float32, эти операции выполняются в float32 и производят выходные данные float32, независимо от того, включен ли автокаст.