torch.autograd.functional.jacobian

torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False, strategy='reverse-mode') [исходный код]

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

Параметры:

• func (function) – функция Python, принимающая тензоры и возвращающая кортеж тензоров или тензор.
• inputs (tuple of Tensors or Tensor) – входные данные функции func.
• create_graph (bool, optional) – Если True, якобиан будет вычислен дифференцируемым способом. Обратите внимание: если strict имеет значение False, результат не может требовать вычисления градиентов или быть не связанным с входными данными. По умолчанию — False.
• strict (bool, optional) – Если True, будет вызвана ошибка, если мы обнаружим, что существует входное значение, от которого не зависит ни один из выходов. Если False, для таких входных данных в качестве якобиана возвращается тензор из нулей, что соответствует ожидаемому математическому значению. По умолчанию — False.
• vectorize (bool, optional) – Эта функция является экспериментальной. Если вам нужен менее экспериментальный и более производительный вариант, рассмотрите возможность использования torch.func.jacrev() или torch.func.jacfwd(). При вычислении якобиана обычно мы вызываем autograd.grad один раз для каждой строки якобиана. Если этот флаг имеет значение True, мы выполняем только один вызов autograd.grad с batched_grad=True, использующий прототипную функцию vmap. Хотя во многих случаях это должно повысить производительность, поскольку функция всё ещё является экспериментальной, возможны резкие падения производительности. Подробнее см. параметр batched_grad функции torch.autograd.grad().
• strategy (str, optional) – Задайте значение "forward-mode" или "reverse-mode", чтобы определить, будет ли якобиан вычисляться прямым или обратным режимом AD. В настоящее время для "forward-mode" требуется vectorized=True. По умолчанию — "reverse-mode". Если у func больше выходов, чем входов, "forward-mode" обычно работает производительнее. В противном случае предпочтительно использовать "reverse-mode".

Возвращает:

Если имеется один вход и один выход, возвращается один тензор, содержащий якобиан для линеаризованных входных и выходных данных. Если один из них является кортежем, якобиан будет кортежем тензоров. Если оба являются кортежами, якобиан будет кортежем кортежей тензоров, где Jacobian[i][j] содержит якобиан i-го выхода и j-го входа; его размер будет равен конкатенации размеров соответствующего выхода и соответствующего входа, а тип данных и устройство будут совпадать с типом данных и устройством соответствующего входа. Если strategy имеет значение forward-mode, тип данных будет соответствовать выходу, в противном случае — входу.

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

Якобиан (Tensor или вложенный кортеж тензоров)

Пример

>>> def exp_reducer(x):
...     return x.exp().sum(dim=1)
>>> inputs = torch.rand(2, 2)
>>> jacobian(exp_reducer, inputs)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]])

>>> jacobian(exp_reducer, inputs, create_graph=True)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]], grad_fn=<ViewBackward>)

>>> def exp_adder(x, y):
...     return 2 * x.exp() + 3 * y
>>> inputs = (torch.rand(2), torch.rand(2))
>>> jacobian(exp_adder, inputs)
(tensor([[2.8052, 0.0000],
        [0.0000, 3.3963]]),
 tensor([[3., 0.],
         [0., 3.]]))

>>> def linear_model(x):
...     W = torch.tensor([[2.0, -1.0], [0.0, 1.0]])
...     b = torch.tensor([1.0, 0.5])
...     return x @ W.T + b

>>> x = torch.randn(4, 2, requires_grad=True)
>>> jac = jacobian(linear_model, x, vectorize=True)
>>> jac.shape
torch.Size([4, 2, 4, 2])