Spec-Zone.ru › PyTorch 1

Автоматическое дифференцирование - torch.autograd

torch.autograd предоставляет классы и функции для реализации автоматического дифференцирования произвольных скалярных функций. Требуются минимальные изменения в существующем коде - вам нужно только объявить Tensor для вычисления градиентов с ключевым словом requires_grad=True. На данный момент мы поддерживаем автоградиент только для чисел с плавающей точкой Tensor типов (half, float, double и bfloat16) и комплексных Tensor типов (cfloat, cdouble).

backward

Вычисляет сумму градиентов заданных тензоров относительно листьев графа.

grad

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

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

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

Этот API находится в стадии бета-тестирования. Несмотря на то, что сигнатуры функций вряд ли изменятся, улучшение охвата операторов планируется перед тем, как мы будем считать его стабильным.

Для подробных инструкций по использованию этого API см. руководство по forward-mode AD.

forward_ad.dual_level

Менеджер контекста, который включает forward AD.

forward_ad.make_dual

Связывает значение тензора с градиентом прямого распространения, касательной, для создания "дуального тензора", который используется для вычисления градиентов forward AD.

forward_ad.unpack_dual

Распаковывает "дуальный тензор", чтобы получить как его значение Tensor, так и градиент forward AD.

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

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

Этот API находится в стадии бета-тестирования. Несмотря на то, что сигнатуры функций вряд ли изменятся, существенные улучшения производительности планируются перед тем, как мы будем считать его стабильным.

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

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

functional.jacobian

Функция, которая вычисляет якобиан заданной функции.

functional.hessian

Функция, которая вычисляет гессиан заданной скалярной функции.

functional.vjp

Функция, которая вычисляет скалярное произведение между вектором v и якобианом заданной функции в точке, заданной входами.

functional.jvp

Функция, которая вычисляет скалярное произведение между якобианом заданной функции в точке, заданной входами, и вектором v.

functional.vhp

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

functional.hvp

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

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

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

Значения градиента по умолчанию

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

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

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

Если у param уже есть атрибут неразреженного .grad:

  1. Если create_graph=False, backward() накапливается в .grad на месте, что сохраняет её шаги.
  2. Если create_graph=True, backward() заменяет .grad новым тензором .grad + new grad, который пытается (но не гарантирует) соответствие предопределённым шагам .grad.

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

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

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

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

Ручные схемы градиентов

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

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

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

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

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

Variable (устарело)

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

API Variable устарел: переменные больше не нужны для использования автоградиента с тензорами. Автоградиент автоматически поддерживает тензоры с requires_grad установленным в True . Ниже приведён краткий обзор изменений:

  • Variable(tensor) и Variable(tensor, requires_grad) по-прежнему работают как ожидалось, но они возвращают тензоры вместо переменных.
  • var.data — то же самое, что tensor.data.
  • Методы, такие как var.backward(), var.detach(), var.register_hook(), теперь работают с тензорами с теми же именами методов.

Кроме того, теперь можно создавать тензоры с requires_grad=True с помощью фабричных методов, таких как torch.randn(), torch.zeros(), torch.ones(), и другие, как показано ниже:

autograd_tensor = torch.randn((2, 3, 4), requires_grad=True)

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

torch.Tensor.grad

Это свойство по умолчанию None и становится тензором в первый раз, когда вызов backward() вычисляет градиенты для self.

torch.Tensor.requires_grad

Является True , если для этого тензора необходимо вычислить градиенты, False в противном случае.

torch.Tensor.is_leaf

Все тензоры, которые имеют requires_grad, что является False по соглашению, будут листвовыми тензорами.

torch.Tensor.backward([gradient, ...])

Вычисляет градиент текущего тензора по отношению к.

torch.Tensor.detach

Возвращает новый тензор, отсоединенный от текущей графы.

torch.Tensor.detach_

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

torch.Tensor.register_hook(hook)

Регистрирует обратное обратное хук.

torch.Tensor.retain_grad()

Разрешает этому тензору иметь их grad заполненными во время backward().

Функция

class torch.autograd.Function(*args, **kwargs) [source]

Базовый класс для создания пользовательских 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.forward

Выполняет операцию.

Function.backward

Определяет формулу для дифференцирования операции с использованием автоматического дифференцирования в обратном режиме (псевдоним для функции vjp).

Function.jvp

Определяет формулу для дифференцирования операции с использованием автоматического дифференцирования в прямом режиме.

Смеси методов контекста

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

function.FunctionCtx.mark_dirty

Помечает указанные тензоры как измененные в операциях на месте.

function.FunctionCtx.mark_non_differentiable

Помечает выходы как недифференцируемые.

function.FunctionCtx.save_for_backward

Сохраняет указанные тензоры для последующего вызова backward().

function.FunctionCtx.set_materialize_grads

Устанавливает, нужно ли материализовывать тензоры градиента вывода.

Проверка численного градиента

gradcheck

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

gradgradcheck

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

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

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

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

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

Параметры:
  • enabled (bool, необязательно) – Установка этого значения в False делает этот менеджер контекста бесполезным.
  • use_cuda (bool, необязательно) – Включает измерение времени событий CUDA с помощью API cudaEvent. Добавляет примерно 4 мкс накладных расходов на каждую операцию с тензором.
  • record_shapes (bool, необязательно) – Если включено протоколирование форм, информация о размерностях входных данных будет собираться. Это позволяет увидеть, какие размерности использовались под капотом, и далее группировать их с помощью prof.key_averages(group_by_input_shape=True). Обратите внимание, что протоколирование форм может исказить данные профилирования. Рекомендуется проводить отдельные запуски с включенным и выключенным протоколированием форм, чтобы проверить временные характеристики. Вероятнее всего, искажение будет незначительным для самых нижних событий (в случае вложенных вызовов функций). Но для функций более высокого уровня общее время работы процессора может быть искусственно увеличено из-за сбора информации о форме.
  • with_flops (bool, необязательно) – Если with_flops установлено, профилировщик оценит значение FLOPs (операций с плавающей запятой) с использованием формы входных данных оператора. Это позволяет оценить производительность оборудования. В настоящее время этот параметр работает только для операторов матричного умножения и 2D свёртки.
  • profile_memory (bool, необязательно) – отслеживает выделение/освобождение памяти тензоров.
  • with_stack (bool, необязательно) – записывает информацию о коде (файл и номер строки) для операций.
  • with_modules (bool) – записывает иерархию модулей (включая имена функций), соответствующую стеку вызовов операции. Например, если вызов метода forward модуля A включает вызов метода forward модуля B, который содержит оператор aten::add, то иерархия модулей оператора aten::add будет A.B. Обратите внимание, что в настоящее время эта поддержка существует только для TorchScript моделей, а не для моделей в режиме eager.
  • use_kineto (bool, необязательно) – экспериментально, включение профилирования с помощью профилировщика Kineto.
  • use_cpu (bool, необязательно) – профилировать события процессора; установка в False требует use_kineto=True и может использоваться для снижения накладных расходов для профилирования только графического процессора.
  • experimental_config (_ExperimentalConfig) – набор экспериментальных параметров, используемых библиотеками профилирования, такими как Kineto. Обратите внимание, что обратная совместимость не гарантируется.

Пример

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

profiler.profile.export_chrome_trace

Экспортирует список событий как файл трассировки Chrome.

profiler.profile.key_averages

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

profiler.profile.self_cpu_time_total

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

profiler.profile.total_average

Среднее значение всех событий.

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 — это счетчик, увеличивающийся каждый раз при создании нового объекта функции обратного прохода и его сохранении.

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

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

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

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

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

Пример

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

profiler.load_nvprof

Открывает файл трассировки nvprof и анализирует аннотации autograd.

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

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 будет включать или отключать обнаружение аномалий автографа, в зависимости от своего аргумента mode. Он может использоваться как менеджер контекста, так и как функция.

См. detect_anomaly выше для получения подробной информации о поведении обнаружения аномалий.

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

Hooks по умолчанию для сохраненных тензоров

Некоторые операции требуют сохранения промежуточных результатов во время прямой передачи для выполнения обратного прохода. Вы можете определить, как эти сохраненные тензоры должны быть упакованы/распакованы, используя хуки. Общее применение — обменять вычисления на память, сохраняя эти промежуточные результаты на диск или в ЦП, а не оставляя их на графическом процессоре. Это особенно полезно, если вы заметили, что ваша модель помещается в видеопамять во время оценки, но не во время обучения. Также см. Hooks для сохраненных тензоров.

class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook) [source]

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

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

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

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

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

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

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

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

class torch.autograd.graph.save_on_cpu(pin_memory=False) [source]

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

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

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

Parameters:

pin_memory (bool) – Если True тензоры будут сохранены в закрепленной памяти ЦП во время упаковки и скопированы на графический процессор асинхронно во время распаковки. По умолчанию 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) [source]

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

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

Parameters:

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

Пример:

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

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/1.13/autograd.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API