Пакет автоматического дифференцирования — torch.autograd

Автоматическое дифференцирование в прямом режиме

Подробные инструкции по использованию этого API см. в руководстве по AD в прямом режиме.

Функциональный API высокого уровня

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

Этот API работает с функциями, предоставленными пользователем, которые принимают на вход только тензоры и возвращают только тензоры. Если ваша функция принимает другие аргументы, не являющиеся тензорами, или тензоры, для которых не задано requires_grad, их можно захватить с помощью лямбда-функции. Например, функцию f с тремя входными аргументами: тензором, для которого требуется вычислить якобиан, другим тензором, который следует считать константой, и логическим флагом в качестве f(input, constant, flag=flag), можно использовать следующим образом: functional.jacobian(lambda x: f(x, constant, flag=flag), input).

Локальное отключение вычисления градиентов

Дополнительную информацию о различиях между режимами no-grad и inference, а также о других связанных механизмах, которые можно спутать с этими двумя режимами, см. в разделе Локальное отключение вычисления градиентов. Список функций для локального отключения градиентов см. также в разделе Локальное отключение вычисления градиентов.

Расположение градиентов по умолчанию

Если нес разреженный param получает нес разреженный градиент при вызове torch.autograd.backward() или torch.Tensor.backward(), param.grad накапливается следующим образом.

Если param.grad изначально имеет значение None:

1. Если память param не перекрывается и является плотной, .grad создается с шагами, соответствующими param (то есть совпадающими с расположением param).
2. В противном случае .grad создается с непрерывными построчными шагами.

Если param уже имеет нес разреженный атрибут .grad:

3. Если create_graph=False, backward() накапливается в .grad на месте, что сохраняет его шаги.
4. Если create_graph=True, backward() заменяет .grad новым тензором .grad + new grad, который пытается (но не гарантирует) сохранить шаги уже существующего .grad.

Для наилучшей производительности рекомендуется использовать поведение по умолчанию (оставлять .grads со значением None перед первым backward(), чтобы их расположение создавалось согласно пунктам 1 или 2 и сохранялось со временем согласно пунктам 3 или 4). Вызовы model.zero_grad() или optimizer.zero_grad() не влияют на расположение .grad.

На самом деле, сброс всех .grads в None перед каждым этапом накопления, например:

for iterations...
    ...
    for param in model.parameters():
        param.grad = None
    loss.backward()

чтобы каждый раз создавать их заново согласно пунктам 1 или 2, — допустимая альтернатива model.zero_grad() или optimizer.zero_grad(), которая может повысить производительность некоторых сетей.

Операции над тензорами на месте

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

Variable (устаревший интерфейс)

Функции autograd для тензоров

Функция

class torch.autograd.Function(*args, **kwargs) [исходный код]

Базовый класс для создания пользовательских autograd.Function.

Чтобы создать пользовательскую autograd.Function, унаследуйте этот класс и реализуйте статические методы forward() и backward(). Затем, чтобы использовать пользовательскую операцию в прямом проходе, вызовите метод класса apply. Не вызывайте forward() напрямую.

Для обеспечения корректности и наилучшей производительности вызывайте правильные методы для ctx и проверяйте функцию обратного прохода с помощью torch.autograd.gradcheck().

Дополнительные сведения об использовании этого класса см. в разделе Расширение torch.autograd.

Примеры:

>>> class Exp(Function):
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result
>>>
>>> # Use it by calling the apply method:
>>> output = Exp.apply(input)

Примеси методов контекста

При создании нового Function доступны следующие методы для ctx.

class torch.autograd.function.FunctionCtx [исходный код]

class torch.autograd.function.FunctionMeta(name, bases, attrs) [исходный код]

Метакласс Function.

Этот метакласс задает следующие свойства:

_backward_cls: класс Function, соответствующий дифференцированной

версии этой функции (создается этим метаклассом на лету).

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

Декоратор для метода обратного прохода.

Базовая пользовательская функция Function, используемая для создания утилит PyTorch

Численная проверка градиента

Профилировщик

Autograd включает профилировщик, позволяющий оценить затраты на выполнение различных операторов внутри модели — как на CPU, так и на GPU. На данный момент реализовано три режима: только CPU с использованием profile; на основе nvprof (регистрирует активность как CPU, так и GPU) с использованием emit_nvtx; и на основе профилировщика vtune с использованием emit_itt.

class torch.autograd.profiler.profile(enabled=True, *, use_device=None, record_shapes=False, with_flops=False, profile_memory=False, with_stack=False, with_modules=False, use_kineto=False, use_cpu=True, experimental_config=None, acc_events=False, custom_trace_id_callback=None, post_processing_timeout_s=None, activity_filters=None) [source]

Менеджер контекста, который управляет состоянием профилировщика autograd и хранит сводку результатов.

Примечание

Это низкоуровневая реализация; большинству пользователей следует использовать вместо неё torch.profiler.

Внутри он просто записывает события выполняемых функций на C++ и предоставляет эти события Python. В него можно поместить любой код, и он будет сообщать только время выполнения функций PyTorch. Примечание: профилировщик привязан к потоку и автоматически передаётся асинхронным задачам.

Параметры:

• enabled (bool, необязательно) – Если установить значение False, менеджер контекста ничего не будет делать.
• use_device (str, необязательно) – Включает измерение времени событий устройства. При использовании CUDA добавляет около 4 мкс накладных расходов на каждую операцию с тензором. Допустимые значения параметра устройства: ‘cuda’, ‘xpu’, ‘mtia’ и ‘privateuseone’.
• record_shapes (bool, необязательно) – Если включена запись форм, будет собираться информация о размерностях входных данных. Это позволяет узнать, какие размерности использовались внутри системы, а затем сгруппировать данные по ним с помощью prof.key_averages(group_by_input_shape=True). Обратите внимание, что запись форм может исказить данные профилирования. Для проверки времени рекомендуется выполнять отдельные запуски с записью форм и без неё. Скорее всего, искажение будет незначительным для событий самого нижнего уровня (при вложенных вызовах функций). Однако для функций более высокого уровня общее собственное время CPU может быть искусственно увеличено из-за сбора информации о формах.
• with_flops (bool, необязательно) – Если включён параметр with_flops, профилировщик оценит количество FLOP (операций с плавающей точкой) на основе формы входных данных оператора. Это позволяет оценить производительность оборудования. В настоящее время этот параметр работает только для операторов умножения матриц и двумерной свёртки.
• profile_memory (bool, необязательно) – отслеживать выделение и освобождение памяти тензоров.
• with_stack (bool, необязательно) – записывать сведения об исходном коде (имя файла и номер строки) для операций.

• with_modules (bool) –

  записывать иерархию модулей (включая имена функций), соответствующую стеку вызовов операции. Например, если прямой вызов модуля A вызывает прямой вызов модуля B, содержащий операцию aten::add, то иерархия модулей для aten::add — A.B.

  Устарело с момента появления параметра ``with_modules``: параметр устарел и будет удалён в будущей версии. Он собирает данные только для моделей TorchScript, которые сами устарели, и ничего не делает в eager-режиме. Используйте with_stack=True, который записывает события nn.Module для моделей в eager-режиме.

• use_kineto (bool, необязательно) – экспериментальный параметр; включает профилирование с помощью профилировщика Kineto.
• use_cpu (bool, необязательно) – профилировать события CPU; значение False требует use_kineto=True и может использоваться для снижения накладных расходов при профилировании только GPU.
• experimental_config (_ExperimentalConfig) – Набор экспериментальных параметров, используемых библиотеками профилирования, такими как Kineto. Обратите внимание: обратная совместимость не гарантируется.
• acc_events (bool) – включить накопление FunctionEvents в нескольких циклах профилирования.
• post_processing_timeout_s (float) – Необязательный тайм-аут постобработки результатов профилирования в секундах. В данном контексте постобработка выполняется после завершения самого профилирования. Если задано значение, разбор событий прекратится по истечении этого времени, и будут возвращены частичные результаты. Полезно для обработки больших трассировок, разбор которых может занять слишком много времени.

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

Включение профилирования памяти или атрибуции исходного кода увеличивает накладные расходы профилировщика.

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

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

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

Из-за некоторых ограничений многопроцессорной обработки CUDA (см. раздел CUDA в многопроцессорной обработке) нельзя использовать профилировщик с use_device = 'cuda' для измерения производительности DataLoader с num_workers > 0. Чтобы измерить производительность загрузки данных, используйте use_device = None или num_workers = 0.

Пример

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
>>>     for _ in range(100):  # any normal python code, really!
>>>         y = x ** 2
>>>         y.backward()
>>> # NOTE: some columns were removed for brevity
>>> print(prof.key_averages().table(sort_by="self_cpu_time_total"))
-----------------------------------  ---------------  ---------------  ---------------
Name                                 Self CPU total   CPU time avg     Number of Calls
-----------------------------------  ---------------  ---------------  ---------------
mul                                  32.048ms         32.048ms         200
pow                                  27.041ms         27.041ms         200
PowBackward0                         9.727ms          55.483ms         100
torch::autograd::AccumulateGrad      9.148ms          9.148ms          100
torch::autograd::GraphRoot           691.816us        691.816us        100
-----------------------------------  ---------------  ---------------  ---------------

class torch.autograd.profiler.emit_nvtx(enabled=True, record_shapes=False) [source]

Менеджер контекста, который заставляет каждую операцию autograd создавать диапазон NVTX.

Полезен при запуске программы под nvprof:

nvprof --profile-from-start off -o trace_name.prof -- <regular command here>

К сожалению, заставить nvprof сбросить собранные данные на диск невозможно, поэтому при профилировании CUDA нужно использовать этот менеджер контекста для аннотирования трассировок nvprof и дождаться завершения процесса, прежде чем изучать их. Затем для визуализации временной шкалы можно использовать NVIDIA Visual Profiler (nvvp), либо результаты можно загрузить с помощью torch.autograd.profiler.load_nvprof(), например, для просмотра в интерактивной оболочке Python.

Параметры:

• enabled (bool, необязательно) – Если задано enabled=False, менеджер контекста ничего не будет делать. Значение по умолчанию: True.
• record_shapes (bool, необязательно) – Если задано record_shapes=True, диапазон nvtx, обрамляющий каждую операцию autograd, будет дополнен информацией о размерах аргументов Tensor, полученных этой операцией, в следующем формате: [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] Аргументы, не являющиеся тензорами, будут представлены как []. Аргументы перечисляются в том порядке, в котором их получает операция бэкенда. Обратите внимание, что этот порядок может отличаться от порядка передачи аргументов на стороне Python. Также учтите, что запись форм может увеличить накладные расходы на создание диапазонов nvtx. Значение по умолчанию: False

Пример

>>> with torch.cuda.profiler.profile():
...     model(x)  # Warmup CUDA memory allocator and profiler
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

Корреляция прямого и обратного проходов

При просмотре профиля, созданного с помощью emit_nvtx в Nvidia Visual Profiler, бывает сложно сопоставить каждую операцию обратного прохода с соответствующей операцией прямого прохода. Чтобы упростить эту задачу, emit_nvtx добавляет к создаваемым диапазонам информацию о номерах последовательности.

Во время прямого прохода каждый диапазон функции помечается seq=<N>. seq — это монотонно увеличивающийся счётчик, который увеличивается при создании каждого нового объекта Function обратного прохода и сохранении его для обратного прохода. Таким образом, аннотация seq=<N>, связанная с каждым диапазоном функции прямого прохода, указывает, что если эта функция создаст объект Function обратного прохода, этому объекту будет присвоен номер последовательности N. Во время обратного прохода диапазон верхнего уровня, охватывающий вызов apply() каждой функции обратного прохода C++, помечается stashed seq=<M>. M — это номер последовательности, присвоенный объекту обратного прохода при создании. Сравнив номера stashed seq в обратном проходе с номерами seq в прямом проходе, можно определить, какая операция прямого прохода создала каждый объект Function обратного прохода.

Функции, выполняемые во время обратного прохода, также помечаются seq=<N>. При стандартном обратном проходе (с create_graph=False) эта информация не имеет значения; фактически для всех таких функций N может быть равен 0. Полезны только диапазоны верхнего уровня, связанные с методами apply() объектов Function обратного прохода: они позволяют сопоставить эти объекты Function с предшествующим прямым проходом.

Двойной обратный проход

Если же выполняется обратный проход с create_graph=True (другими словами, если настраивается двойной обратный проход), выполнение каждой функции в ходе обратного прохода получает ненулевой полезный seq=<N>. Эти функции сами могут создавать объекты Function для последующего выполнения во время двойного обратного прохода, как это делают исходные функции прямого прохода. Связь между обратным и двойным обратным проходами концептуально такая же, как между прямым и обратным проходами: функции по-прежнему создают диапазоны с меткой текущего номера последовательности, созданные ими объекты Function по-прежнему сохраняют эти номера последовательности, а во время последующего двойного обратного прохода диапазоны apply() объектов Function по-прежнему помечаются номерами stashed seq, которые можно сравнить с номерами seq из обратного прохода.

class torch.autograd.profiler.emit_itt(enabled=True, record_shapes=False) [source]

Менеджер контекста, который заставляет каждую операцию autograd создавать диапазон ITT.

Полезен при запуске программы под Intel(R) VTune Profiler:

vtune <--vtune-flags> <regular command here>

API Instrumentation and Tracing Technology (ITT) позволяет приложению создавать трассировочные данные и управлять их сбором во время выполнения в различных инструментах Intel. Этот менеджер контекста предназначен для аннотирования трассировки профилирования Intel(R) VTune. С его помощью в графическом интерфейсе Intel(R) VTune Profiler можно просматривать помеченные диапазоны.

Параметры:

• enabled (bool, необязательно) – Если задано enabled=False, менеджер контекста ничего не будет делать. Значение по умолчанию: True.
• record_shapes (bool, необязательно) – Если задано record_shapes=True, диапазон itt, обрамляющий каждую операцию autograd, будет дополнен информацией о размерах аргументов Tensor, полученных этой операцией, в следующем формате: [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] Аргументы, не являющиеся тензорами, будут представлены как []. Аргументы перечисляются в том порядке, в котором их получает операция бэкенда. Обратите внимание, что этот порядок может отличаться от порядка передачи аргументов на стороне Python. Также учтите, что запись форм может увеличить накладные расходы на создание диапазонов itt. Значение по умолчанию: False

Пример

>>> with torch.autograd.profiler.emit_itt():
...     model(x)

Отладка и обнаружение аномалий

class torch.autograd.detect_anomaly(check_nan=True) [source]

Менеджер контекста, включающий обнаружение аномалий в движке autograd.

Он выполняет две задачи:

• Если включить обнаружение во время прямого прохода, при обратном проходе будет выведена трассировка стека операции прямого прохода, создавшей функцию обратного прохода, в которой произошёл сбой.
• Если check_nan равно True, любые вычисления обратного прохода, порождающие значение «nan», вызовут ошибку. Значение по умолчанию — True.

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

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

Пример

>>> import torch
>>> from torch import autograd
>>> class MyFunc(autograd.Function):
...     @staticmethod
...     def forward(ctx, inp):
...         return inp.clone()
...
...     @staticmethod
...     def backward(ctx, gO):
...         # Error during the backward pass
...         raise RuntimeError("Some error in backward")
...         return gO.clone()
>>> def run_fn(a):
...     out = MyFunc.apply(a)
...     return out.sum()
>>> inp = torch.rand(10, 10, requires_grad=True)
>>> out = run_fn(inp)
>>> out.backward()
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward
>>> with autograd.detect_anomaly():
...     inp = torch.rand(10, 10, requires_grad=True)
...     out = run_fn(inp)
...     out.backward()
    Traceback of forward call that caused the error:
      File "tmp.py", line 53, in <module>
        out = run_fn(inp)
      File "tmp.py", line 44, in run_fn
        out = MyFunc.apply(a)
    Traceback (most recent call last):
      File "<stdin>", line 4, in <module>
      File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward

class torch.autograd.set_detect_anomaly(mode, check_nan=True) [source]

Менеджер контекста, включающий или отключающий обнаружение аномалий в движке autograd.

set_detect_anomaly включает или отключает обнаружение аномалий autograd в зависимости от аргумента mode. Его можно использовать как менеджер контекста или как функцию.

Подробнее о поведении обнаружения аномалий см. выше, в разделе detect_anomaly.

Параметры:

• mode (bool) – Флаг, указывающий, следует ли включить обнаружение аномалий (True) или отключить его (False).
• check_nan (bool) – Флаг, указывающий, следует ли вызывать ошибку, если обратный проход порождает значение «nan».

Граф Autograd

Autograd предоставляет методы, позволяющие исследовать граф и изменять поведение во время обратного прохода.

Атрибут grad_fn объекта torch.Tensor содержит torch.autograd.graph.Node, если тензор является результатом операции, зарегистрированной autograd (то есть grad_mode включён и хотя бы для одного входного тензора требуются градиенты), и None в противном случае.

Для выполнения обратного прохода некоторым операциям требуется сохранить промежуточные результаты во время прямого прохода. Эти промежуточные результаты сохраняются в виде атрибутов объекта grad_fn, к которым можно получить доступ. Например:

>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
>>> b = a.exp()
>>> print(isinstance(b.grad_fn, torch.autograd.graph.Node))
True
>>> print(dir(b.grad_fn))
['__call__', '__class__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '_raw_saved_result', '_register_hook_dict', '_saved_result', 'metadata', 'name', 'next_functions', 'register_hook', 'register_prehook', 'requires_grad']
>>> print(torch.allclose(b.grad_fn._saved_result, b))
True

Также можно с помощью хуков определить, как следует упаковывать и распаковывать сохранённые тензоры. Часто это позволяет обменять вычислительные ресурсы на память: сохранять промежуточные результаты на диске или в памяти CPU, а не оставлять их на GPU. Это особенно полезно, если при оценке модель помещается в память GPU, а при обучении — нет. При написании pack_hook, хранящего входной тензор, сначала вызовите для него .detach(), чтобы избежать циклической ссылки, если сохранённый тензор является выходом графа; подробности см. в разделе Хуки для сохранённых тензоров.

class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook) [исходный код]

Менеджер контекста, задающий пару хуков упаковки и распаковки для сохранённых тензоров.

Используйте этот менеджер контекста, чтобы определить, как упаковывать промежуточные результаты операции перед сохранением и распаковывать их при извлечении.

В этом контексте функция pack_hook будет вызываться каждый раз, когда операция сохраняет тензор для обратного прохода (сюда входят промежуточные результаты, сохранённые с помощью save_for_backward(), а также результаты, сохранённые операциями, определёнными в PyTorch). Результат pack_hook сохраняется в графе вычислений вместо исходного тензора.

Функция unpack_hook вызывается, когда требуется получить доступ к сохранённому тензору, а именно при выполнении torch.Tensor.backward() или torch.autograd.grad(). Она принимает в качестве аргумента упакованный объект, возвращённый функцией pack_hook, и должна возвращать тензор с тем же содержимым, что и исходный тензор (переданный на вход соответствующей функции pack_hook).

Хуки должны иметь следующие сигнатуры:

pack_hook(tensor: Tensor) -> Any

unpack_hook(Any) -> Tensor

где возвращаемое значение pack_hook является допустимым входным значением для unpack_hook.

Как правило, значения unpack_hook(pack_hook(t)) и t должны совпадать по значению, размеру, типу данных и устройству.

Пример:

>>> def pack_hook(x):
...     print("Packing", x)
...     return x.detach()
>>>
>>> def unpack_hook(x):
...     print("Unpacking", x)
...     return x
>>>
>>> a = torch.ones(5, requires_grad=True)
>>> b = torch.ones(5, requires_grad=True) * 2
>>> with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
...     y = a * b
Packing tensor([1., 1., 1., 1., 1.], requires_grad=True)
Packing tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)
>>> y.sum().backward()
Unpacking tensor([1., 1., 1., 1., 1.], requires_grad=True)
Unpacking tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)

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

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

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

Одновременно разрешена только одна пара хуков. При рекурсивном вложении этого менеджера контекста применяться будет только самая внутренняя пара хуков.

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

Чтобы избежать циклической ссылки, возвращаемое значение pack_hook не должно содержать ссылку на входной тензор. Например, в качестве хука упаковки используйте lambda x: x.detach() вместо lambda x: x.

class torch.autograd.graph.save_on_cpu(pin_memory=False, device_type='cuda') [исходный код]

Менеджер контекста, в котором тензоры, сохранённые во время прямого прохода, помещаются в память CPU, а затем извлекаются для обратного прохода.

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

Используйте этот менеджер контекста, чтобы обменять вычислительные ресурсы на память GPU (например, если модель не помещается в память GPU во время обучения).

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

При pin_memory=True копирование с GPU на CPU при упаковке выполняется асинхронно. Доступ к сохранённым тензорам на CPU (например, через grad_fn._saved_self) до завершения потока CUDA может привести к получению некорректных данных. Если необходимо прочитать эти данные, сначала вызовите torch.cuda.synchronize().

Параметры:

pin_memory (bool) – Если True, тензоры сохраняются в закреплённой памяти CPU при упаковке и асинхронно копируются на GPU как при упаковке, так и при распаковке. По умолчанию — False. См. также раздел Использование буферов закреплённой памяти.

Пример:

>>> a = torch.randn(5, requires_grad=True, device="cuda")
>>> b = torch.randn(5, requires_grad=True, device="cuda")
>>> c = torch.randn(5, requires_grad=True, device="cuda")
>>>
>>> def f(a, b, c):
...     prod_1 = a * b           # a and b are saved on GPU
...     with torch.autograd.graph.save_on_cpu():
...         prod_2 = prod_1 * c  # prod_1 and c are saved on CPU
...     y = prod_2 * a           # prod_2 and a are saved on GPU
...     return y
>>>
>>> y = f(a, b, c)
>>> del a, b, c  # for illustration only
>>> # the content of a, b, and prod_2 are still alive on GPU
>>> # the content of prod_1 and c only live on CPU
>>> y.sum().backward()  # all CPU tensors are moved back to GPU, for backward
>>> # all intermediary tensors are released (deleted) after the call to backward

class torch.autograd.graph.disable_saved_tensors_hooks(error_message) [исходный код]

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

Полезен при создании функции, несовместимой с хуками по умолчанию для сохранённых тензоров.

Параметры:

error_message (str) – Если хуки по умолчанию для сохранённых тензоров используются после отключения, будет вызван RuntimeError с этим сообщением об ошибке.

Тип возвращаемого значения:

Generator[None, None, None]

Пример:

>>> message = "saved tensors default hooks are disabled"
>>> with torch.autograd.graph.disable_saved_tensors_hooks(message):
...     # Raises RuntimeError: saved tensors default hooks are disabled
...     with torch.autograd.graph.save_on_cpu():
...         pass

class torch.autograd.graph.register_multi_grad_hook(tensors, fn, *, mode='all') [исходный код]

Регистрирует хук обратного прохода для нескольких градиентов.

Поддерживаются два режима: "all" и "any".

В режиме "all" хук вызывается после вычисления градиентов относительно каждого тензора в tensors. Если тензор входит в tensors, но не является частью графа, или если тензор не нужен для вычисления градиентов ни для одного из inputs, указанных в текущем вызове .backward() или .grad(), этот тензор игнорируется, и хук не ожидает вычисления его градиента.

После вычисления градиента каждого не проигнорированного тензора вызывается fn с этими градиентами. Для тензоров, градиенты которых не были вычислены, будет передано значение None.

В режиме "any" хук вызывается после вычисления первого градиента относительно тензора из tensors. Этот градиент передаётся хуку в качестве аргумента.

Хук не должен изменять свои аргументы.

Эта функция возвращает дескриптор с методом handle.remove(), который удаляет хук.

Примечание

Дополнительные сведения о том, когда выполняется этот хук и как порядок его выполнения соотносится с другими хуками, см. в разделе Выполнение хуков обратного прохода.

Пример:

>>> import torch
>>>
>>> a = torch.rand(2, 3, requires_grad=True)
>>> b = torch.rand(2, 3, requires_grad=True)
>>> c = a * b
>>> d = a * b
>>>
>>> def fn(grads):
...     print([g is not None for g in grads])
...
>>> torch.autograd.graph.register_multi_grad_hook((a, b, c, d), fn)
>>>
>>> c.sum().backward(retain_graph=True)
[True, True, True, False]
>>> c.sum().backward(inputs=(a,), retain_graph=True)
[True, False, True, False]
>>>

Тип возвращаемого значения:

RemovableHandle

class torch.autograd.graph.node_creation_hook(hook) [исходный код]

Менеджер контекста, регистрирующий хук, вызываемый для каждого созданного в нём узла autograd.

В этом контексте функция hook вызывается один раз для каждого узла графа autograd, созданного операциями над тензорами, для которых требуются градиенты. Единственным аргументом ей передаётся только что созданный torch.autograd.graph.Node. Предполагается, что хук будет записывать узел, добавлять записи в node.metadata или регистрировать для него хуки обратного прохода с помощью register_hook() и register_prehook().

Узел передаётся хуку только после полной настройки: его next_functions связаны, метаданные всех выходов заданы, а тензоры, сохранённые для обратного прохода (_saved_*), сохранены. Поэтому хук, исследующий узел, видит его полностью сформированное состояние.

Хук должен иметь следующую сигнатуру:

hook(node: torch.autograd.graph.Node) -> None

Регистрация действует в пределах потока и распространяется так же, как и другое локальное состояние потока autograd: она активна в рабочих потоках движка autograd, поэтому хук также вызывается для узлов, созданных во время обратного прохода (например, с помощью create_graph=True или при повторном вычислении контрольной точки).

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

Один из вариантов применения — отнесение работы, выполняемой во время обратного прохода, к области прямого прохода, создавшей граф: для этого состояние фиксируется при создании узла и восстанавливается во время выполнения обратного прохода для этого узла:

>>> def creation_hook(node):
...     # ``current_region()``/``enter_region()`` are user-defined and
...     # stand in for whatever thread-local state you want to restore
...     # while this node runs in backward.
...     region = current_region()
...     node.register_prehook(lambda gO: enter_region(region))
...     node.register_hook(lambda gI, gO: enter_region(None))
>>>
>>> with torch.autograd.graph.node_creation_hook(creation_hook):
...     loss = model(inputs)

Пример:

>>> a = torch.ones(5, requires_grad=True)
>>> with torch.autograd.graph.node_creation_hook(lambda node: print(node)):
...     b = a * 2
<AccumulateGrad object at ...>
<MulBackward0 object at ...>

Примечание

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

class torch.autograd.graph.allow_mutation_on_saved_tensors [исходный код]

Менеджер контекста, разрешающий изменять тензоры, сохранённые для обратного прохода.

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

Чтобы обеспечить корректное поведение, прямой и обратный проходы следует выполнять в одном и том же менеджере контекста.

Возвращает:

Объект _AllowMutationOnSavedContext, содержащий состояние, которым управляет этот менеджер контекста. Этот объект может быть полезен для отладки. Состояние, которым управляет менеджер контекста, автоматически очищается при выходе из него.

Тип возвращаемого значения:

Generator[_AllowMutationOnSavedContext, None, None]

Пример:

>>> import torch
>>> with torch.autograd.graph.allow_mutation_on_saved_tensors():
...     # forward
...     a = torch.ones(2, 3, requires_grad=True)
...     b = a.clone()
...     out = (b**2).sum()
...     b.sin_()
...     # backward
...     out.sum().backward()
...
tensor([[0.8415, 0.8415, 0.8415],
        [0.8415, 0.8415, 0.8415]], grad_fn=<SinBackward0>)

class torch.autograd.graph.GradientEdge(node, output_nr, ownership_token=None) [исходный код]

Объект, представляющий заданное ребро градиента в графе autograd.

Чтобы получить ребро градиента, в котором будет вычисляться градиент заданного тензора, выполните edge = autograd.graph.get_gradient_edge(tensor).

torch.autograd.graph.get_gradient_edge(tensor) [исходный код]

Получает ребро градиента для вычисления градиента заданного Tensor.

В частности, это эквивалентно вызову g = autograd.grad(loss, input) и g = autograd.grad(loss, get_gradient_edge(input)).

Тип возвращаемого значения:

GradientEdge

torch.autograd.graph.region_activation_memory_budget(budget) [исходный код]

Менеджер контекста, задающий бюджет памяти для активаций в области скомпилированного прямого прохода, трассируемого внутри этого контекста.

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

Это экспериментальная функция, которая может измениться.

При использовании torch.compile() алгоритм разбиения min-cut выбирает, какие активации сохранить, а какие повторно вычислить в обратном проходе, чтобы не превысить бюджет памяти. budget — это отношение в [0, 1]: 0.0 соответствует объёму памяти для активаций при применении контрольных точек активации ко всей области, а 1.0 — объёму памяти для активаций при использовании стратегии по умолчанию, оптимизированной во время выполнения. Таким образом, 0.4 задаёт стратегию, при которой сохраняется на 40% меньше активаций по сравнению со стратегией по умолчанию. Алгоритм решает задачу о рюкзаке 0-1, чтобы определить минимальный объём повторных вычислений, необходимый для соблюдения бюджета. Это переопределяет глобальное значение torch._functorch.config.activation_memory_budget для аннотированной области.

Примечание

В настоящее время алгоритм разбиения поддерживает только один бюджет на скомпилированный граф, поэтому аннотация должна охватывать все операции прямого прохода в графе (частичная аннотация отклоняется, а не применяется незаметно ко всему графу), и бюджет для всех аннотированных узлов должен совпадать. Чтобы задать разные бюджеты для разных частей модели, разделите их разрывом графа (например, с помощью torch._dynamo.graph_break()), чтобы каждая часть стала отдельным графом.

Эта функция действует только при использовании torch.compile(); её использование вне скомпилированной области приводит к RuntimeError.

Параметры:

budget (float) – Отношение бюджета памяти для активаций в [0, 1].

Тип возвращаемого значения:

AbstractContextManager[None]

Пример:

>>> with torch.autograd.graph.region_activation_memory_budget(0.0):
...     x = layer(x)  # recompute this region's activations in backward

torch.autograd.graph.set_warn_on_accumulate_grad_stream_mismatch(enabled) [исходный код]

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

torch.autograd.graph.set_override_stale_capture_stream(enabled) [исходный код]

Управляет поведением при обнаружении autograd устаревшего потока, не участвующего в захвате, во время захвата графа CUDA.

Во время захвата графа CUDA узлы autograd могут ссылаться на устаревший поток, не входящий в захват. Если флаг отключён (это начальное состояние процесса), autograd вызывает RuntimeError, когда устаревшим является поток по умолчанию (поток 0), поскольку в этом случае захват всегда становится недействительным: cudaStreamWaitEvent в потоке по умолчанию добавляет в граф поток, не участвующий в захвате. Для устаревших потоков, отличных от потока по умолчанию, ссылка на поток не изменяется; захват будет успешным, если пользователь присоединил поток к захвату (например, с помощью capture_stream.wait_stream(stale_stream)), и в противном случае завершится ошибкой среды выполнения CUDA.

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

Флаг также управляет синхронизацией в конце обратного прохода между потоком каждого листового тензора и текущим потоком вызывающего кода. Если все входящие градиенты листового тензора не определены (например, при вычислении некоторых градиентов вне основного процесса с возвратом None из autograd), эту ситуацию нельзя обработать с помощью переопределения. Поэтому, если только один из двух потоков участвует в захвате, синхронизация пересечёт границу захвата: если флаг включён, синхронизация пропускается (работа в потоке, не участвующем в захвате, не входит в захват, поэтому порядок её выполнения не имеет значения); если флаг отключён, вызывается RuntimeError.

Параметры:

enabled (bool) – Если True, во время захвата графа CUDA заменять устаревшие потоки, не участвующие в захвате, захватывающим потоком производителя, а также пропускать синхронизацию листовых тензоров в конце обратного прохода, если она пересекает границу захвата. Если False (начальное состояние процесса), вызывать ошибку, когда устаревшим является поток по умолчанию (поток 0) или когда синхронизация листового тензора пересекает границу захвата; остальные устаревшие потоки остаются без изменений.