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

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

torch.amp.autocast_mode.is_autocast_available(device_type) [источник]

Возвращает логическое значение, указывающее, доступно ли автоматическое приведение типов на device_type.

Параметры:

device_type (str) – Тип устройства. Возможные значения: ‘cuda’, ‘cpu’, ‘mtia’, ‘maia’, ‘xpu’ и т. д. Тип совпадает с атрибутом type объекта torch.device. Таким образом, тип устройства тензора можно получить с помощью Tensor.device.type.

Возвращаемый тип:

bool

class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None) [источник]

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

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

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

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

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

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

class AutocastModel(nn.Module):
    ...

    @torch.autocast(device_type="cuda")
    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 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:

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

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

Параметры:

• device_type (str, обязательный) – Тип устройства. Возможные значения: ‘cuda’, ‘cpu’, ‘mtia’, ‘maia’, ‘xpu’ и ‘hpu’. Тип совпадает с атрибутом type объекта torch.device. Таким образом, тип устройства тензора можно получить с помощью Tensor.device.type.
• enabled (bool, необязательный) – Следует ли включить автоматическое приведение типов в этой области. Значение по умолчанию: True
• dtype (torch_dtype, необязательный) – Тип данных для операций, выполняемых с автоматическим приведением типов. Если dtype равен None, используется значение по умолчанию (torch.float16 для CUDA и torch.bfloat16 для CPU), задаваемое функцией get_autocast_dtype(). Значение по умолчанию: None
• cache_enabled (bool, необязательный) – Следует ли включить кэш весов для автоматического приведения типов. Значение по умолчанию: True

torch.amp.custom_fwd(fwd=None, *, device_type, cast_inputs=None) [источник]

Создаёт вспомогательный декоратор для методов forward пользовательских функций autograd.

Функции autograd являются подклассами torch.autograd.Function. Подробнее см. на странице с примерами.

Параметры:

• device_type (str) – Тип устройства: ‘cuda’, ‘cpu’, ‘mtia’, ‘maia’, ‘xpu’ и т. д. Тип совпадает с атрибутом type объекта torch.device. Таким образом, тип устройства тензора можно получить с помощью Tensor.device.type.
• cast_inputs (torch.dtype или None, необязательный, по умолчанию None) – Если значение не равно None и forward выполняется в области с включённым автоматическим приведением типов, входные тензоры с плавающей точкой преобразуются к целевому типу данных (тензоры без плавающей точки не затрагиваются), после чего forward выполняется с отключённым автоматическим приведением типов. Если значение равно None, внутренние операции forward выполняются в текущем состоянии автоматического приведения типов.

Примечание

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

torch.amp.custom_bwd(bwd=None, *, device_type) [источник]

Создаёт вспомогательный декоратор для методов обратного прохода пользовательских функций autograd.

Функции autograd являются подклассами torch.autograd.Function. Гарантирует, что backward выполняется в том же состоянии автоматического приведения типов, что и forward. Подробнее см. на странице с примерами.

Параметры:

device_type (str) – Тип устройства: ‘cuda’, ‘cpu’, ‘mtia’, ‘maia’, ‘xpu’ и т. д. Тип совпадает с атрибутом type объекта torch.device. Таким образом, тип устройства тензора можно получить с помощью Tensor.device.type.

class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True) [источник]

См. torch.autocast.

torch.cuda.amp.autocast(args...) устарел. Используйте вместо него torch.amp.autocast("cuda", args...).

torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None) [источник]

torch.cuda.amp.custom_fwd(args...) устарел. Используйте вместо него torch.amp.custom_fwd(args..., device_type='cuda').

torch.cuda.amp.custom_bwd(bwd) [источник]

torch.cuda.amp.custom_bwd(args...) устарел. Используйте вместо него torch.amp.custom_bwd(args..., device_type='cuda').

class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True) [источник]

См. torch.autocast. torch.cpu.amp.autocast(args...) устарел. Используйте вместо него torch.amp.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) [источник]

См. torch.amp.GradScaler. torch.cuda.amp.GradScaler(args...) устарел. Используйте вместо него torch.amp.GradScaler("cuda", args...).

class torch.cpu.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True) [источник]

См. torch.amp.GradScaler. torch.cpu.amp.GradScaler(args...) устарел. Используйте вместо него torch.amp.GradScaler("cpu", args...).

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