wehub-resource-sync
754 lines
25 KiB
Python
754 lines
25 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import typing
|
|
from typing import TYPE_CHECKING, TypeVar, overload
|
|
|
|
import paddle
|
|
from paddle.base import framework
|
|
from paddle.incubate.autograd import primapi, utils
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Sequence
|
|
|
|
from paddle import Tensor
|
|
from paddle._typing import TensorOrTensors
|
|
|
|
_OutputT = TypeVar("_OutputT", Tensor, tuple[Tensor, ...])
|
|
|
|
|
|
@overload
|
|
def vjp(
|
|
func: Callable[..., _OutputT],
|
|
xs: Tensor,
|
|
v: TensorOrTensors | None = None,
|
|
) -> tuple[_OutputT, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def vjp(
|
|
func: Callable[..., _OutputT],
|
|
xs: Sequence[Tensor],
|
|
v: TensorOrTensors | None = None,
|
|
) -> tuple[_OutputT, tuple[Tensor, ...]]: ...
|
|
|
|
|
|
def vjp(func, xs, v=None):
|
|
r"""Computes the Vector-Jacobian product, a functional form of
|
|
reverse mode automatic differentiation.
|
|
|
|
Warning:
|
|
This API is in beta, the signatures could be changed in future version.
|
|
|
|
Args:
|
|
func(Callable): A function that takes ``xs`` as inputs parameter and
|
|
returns a sequence of Tensors or a Tensor.
|
|
xs(Tensor|Sequence[Tensor]): Used as positional arguments to evaluate
|
|
``func``. ``xs`` is accepted as one Tensor or a sequence of Tensors.
|
|
v(Tensor|Sequence[Tensor]|None, optional): The cotangent vector involved
|
|
in the VJP computation. ``v`` matches the size and shape of
|
|
``func`` 's output. Defaults to None, which is equivalent to all
|
|
ones the same size of ``func`` 's output.
|
|
|
|
Returns:
|
|
output(tuple):
|
|
|
|
- func_out(Tensor|tuple[Tensor]): The output of ``func(xs)`` .
|
|
- vjp(Tensor|tuple[Tensor]): The vjp result.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> def func(x):
|
|
... return paddle.matmul(x, x)
|
|
>>> x = paddle.ones(shape=[2, 2], dtype='float32')
|
|
>>> _, vjp_result = paddle.incubate.autograd.vjp(func, x)
|
|
>>> print(vjp_result)
|
|
Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
|
[[4., 4.],
|
|
[4., 4.]])
|
|
|
|
>>> v = paddle.to_tensor([[1.0, 0.0], [0.0, 0.0]])
|
|
>>> _, vjp_result = paddle.incubate.autograd.vjp(func, x, v)
|
|
>>> print(vjp_result)
|
|
Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
|
[[2., 1.],
|
|
[1., 0.]])
|
|
"""
|
|
_check_inputs(func, xs, v)
|
|
|
|
# ``_separate`` breaks the dependencies between ``xs`` and other
|
|
# variables. See more ``_separate`` .
|
|
if framework.in_dygraph_mode() or not utils.prim_enabled():
|
|
xs, v = _separate(xs), _separate(v)
|
|
ys = func(*xs) if isinstance(xs, typing.Sequence) else func(xs)
|
|
_check_v_shape(v, ys)
|
|
|
|
return ys, _grad(ys, xs, v)
|
|
|
|
|
|
@overload
|
|
def jvp(
|
|
func: Callable[..., _OutputT],
|
|
xs: Tensor,
|
|
v: TensorOrTensors | None = None,
|
|
) -> tuple[_OutputT, Tensor]: ...
|
|
|
|
|
|
@overload
|
|
def jvp(
|
|
func: Callable[..., _OutputT],
|
|
xs: Sequence[Tensor],
|
|
v: TensorOrTensors | None = None,
|
|
) -> tuple[_OutputT, tuple[Tensor, ...]]: ...
|
|
|
|
|
|
def jvp(func, xs, v=None):
|
|
r"""
|
|
Computes the Jacobian-Vector product for a function at the given
|
|
inputs and a vector in the tangent space induced by the inputs.
|
|
|
|
Warning:
|
|
This API is in beta, the signatures could be changed in future version.
|
|
|
|
Args:
|
|
func(Callable): The ``func`` takes as input a Tensor or a Sequence
|
|
of Tensors and returns a Tensor or a Sequence of Tensors.
|
|
xs(Tensor|Sequence[Tensor]): Used as positional arguments to
|
|
evaluate ``func``. The ``xs`` is accepted as one Tensor or a
|
|
Sequence of Tensors.
|
|
v(Tensor|Sequence[Tensor]|None, Optional): The tangent vector involved
|
|
in the JVP computation. The ``v`` matches the size and shape of
|
|
``xs`` . Default value is None and in this case is equivalent to
|
|
all ones the same size of ``xs`` .
|
|
|
|
Returns:
|
|
output(tuple):
|
|
|
|
- func_out(Tensor|tuple[Tensor]): The output of ``func(xs)`` .
|
|
- jvp(Tensor|tuple[Tensor]): The jvp result.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> def func(x):
|
|
... return paddle.matmul(x, x)
|
|
>>> x = paddle.ones(shape=[2, 2], dtype='float32')
|
|
>>> _, jvp_result = paddle.incubate.autograd.jvp(func, x)
|
|
>>> print(jvp_result)
|
|
Tensor(shape=[2, 2], dtype=float32, place=Place(gpu:0), stop_gradient=False,
|
|
[[4., 4.],
|
|
[4., 4.]])
|
|
|
|
>>> v = paddle.to_tensor([[1.0, 0.0], [0.0, 0.0]])
|
|
>>> _, jvp_result = paddle.incubate.autograd.jvp(func, x, v)
|
|
>>> print(jvp_result)
|
|
Tensor(shape=[2, 2], dtype=float32, place=Place(gpu:0), stop_gradient=False,
|
|
[[2., 1.],
|
|
[1., 0.]])
|
|
|
|
"""
|
|
_check_inputs(func, xs, v)
|
|
# ``_separate`` breaks the dependencies between ``xs`` and other
|
|
# variables. See more ``_separate`` .
|
|
if framework.in_dygraph_mode() or not utils.prim_enabled():
|
|
xs, v = _separate(xs), _separate(v)
|
|
ys = func(*xs) if isinstance(xs, typing.Sequence) else func(xs)
|
|
_check_v_shape(v, xs)
|
|
|
|
if not framework.in_dygraph_mode() and utils.prim_enabled():
|
|
return ys, primapi.forward_grad(ys, xs, v)
|
|
else:
|
|
return ys, _double_backward_trick(ys, xs, v)
|
|
|
|
|
|
def _double_backward_trick(ys, xs, v):
|
|
"""Double backward trick for computing ``jvp`` by ``vjp``
|
|
see details: https://j-towns.github.io/2017/06/12/A-new-trick.html
|
|
"""
|
|
# The value of ys_grad is not important, it can be any random value in
|
|
# theory, but it's required to set stop_gradient=False.
|
|
ys_grad = _zeros_like_with_grad(ys)
|
|
xs_grad = _grad(ys, xs, ys_grad)
|
|
return _grad(xs_grad, ys_grad, v)
|
|
|
|
|
|
def _zeros_like_with_grad(xs):
|
|
"""Create a zero or zeros sequence Tensor like ``xs`` with a flag
|
|
``stop_gradient=False`` .
|
|
"""
|
|
if not isinstance(xs, typing.Sequence):
|
|
ys = paddle.zeros_like(xs)
|
|
ys.stop_gradient = False
|
|
else:
|
|
ys = []
|
|
for x in xs:
|
|
y = paddle.zeros_like(x)
|
|
y.stop_gradient = False
|
|
ys.append(y)
|
|
return ys
|
|
|
|
|
|
class Jacobian:
|
|
r"""
|
|
Computes the Jacobian matrix of a given function.
|
|
|
|
If the function has multiple inputs and multiple outputs, during internal
|
|
implementation, all input tensors are concatenated after being flatten,
|
|
the batch dimension is retained, and the output is subject to the same
|
|
processing rules.
|
|
|
|
Once the Jacobian ``J`` is constructed, you can use a multidimensional index
|
|
to retrieve the submatrix of ``J``, as same as slicing a Tensor. The
|
|
submatrix is lazily evaluated along row axis, and will be cached once
|
|
evaluated.
|
|
|
|
For examples, supposing ``is_batched=True``, you can retrieve the submatrix
|
|
by following methods:
|
|
|
|
* J[:], retrieving the full matrix.
|
|
* J[:, :, j], retrieving the partial derivatives w.r.t. the j'th input
|
|
variable.
|
|
* J[:, i, :], retrieving the partial derivatives w.r.t. the i'th output
|
|
variable.
|
|
* J[:, i, j], retrieving the partial derivatives w.r.t. the i'th output
|
|
variable and the j'th input variable.
|
|
|
|
Notes:
|
|
|
|
Ellipsis index is not supported currently.
|
|
|
|
Warning:
|
|
This API is in beta, the signatures could be changed in future version.
|
|
|
|
Args:
|
|
|
|
func (Callable): A python function that takes a Tensor or a sequence of
|
|
Tensors as inputs(the first dimension is batch size) and
|
|
returns a Tensor a sequence of Tensors.
|
|
xs (Tensor|Sequence[Tensor]): The input to the function ``func`` .
|
|
is_batched (bool): If true, the first axis is batch axis. Defaults to
|
|
False.
|
|
|
|
Returns:
|
|
|
|
Jacobian (Object): A python object retains the Jacobian matrix.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> def func(x, y):
|
|
... return paddle.matmul(x, y)
|
|
>>> x = paddle.to_tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
>>> J = paddle.incubate.autograd.Jacobian(func, [x, x])
|
|
>>> print(J[:, :])
|
|
Tensor(shape=[4, 8], dtype=float32, place=Place(cpu), stop_gradient=False,
|
|
[[1., 3., 0., 0., 1., 0., 2., 0.],
|
|
[2., 4., 0., 0., 0., 1., 0., 2.],
|
|
[0., 0., 1., 3., 3., 0., 4., 0.],
|
|
[0., 0., 2., 4., 0., 3., 0., 4.]])
|
|
|
|
>>> print(J[0, :])
|
|
Tensor(shape=[8], dtype=float32, place=Place(cpu), stop_gradient=False,
|
|
[1., 3., 0., 0., 1., 0., 2., 0.])
|
|
>>> print(J[:, 0])
|
|
Tensor(shape=[4], dtype=float32, place=Place(cpu), stop_gradient=False,
|
|
[1., 2., 0., 0.])
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
func: Callable[..., TensorOrTensors],
|
|
xs: TensorOrTensors,
|
|
is_batched: bool = False,
|
|
) -> None:
|
|
if not is_batched:
|
|
self._jacobian = _JacobianNoBatch(func, xs)
|
|
else:
|
|
self._jacobian = _JacobianBatchFirst(func, xs)
|
|
|
|
def __getitem__(
|
|
self, indexes: int | slice | tuple[int | slice, ...]
|
|
) -> Tensor:
|
|
return self._jacobian[indexes]
|
|
|
|
@property
|
|
def shape(self) -> list[int]:
|
|
"""The shape of flattened Jacobian matrix."""
|
|
return self._jacobian.shape
|
|
|
|
|
|
class Hessian:
|
|
"""
|
|
Computes the Hessian matrix with a given ``func`` with respect to ``xs`` .
|
|
|
|
If the function has multiple inputs, during internal implementation,
|
|
all input tensors are concatenated after being flatten, the batch dimension
|
|
is retained.
|
|
|
|
The Hessian submatrix is lazily evaluated, and can be retrieved with a
|
|
multidimensional indexes. See details ``Jacobian`` .
|
|
|
|
Warning:
|
|
This API is in beta, the signatures could be changed in future version.
|
|
|
|
Args:
|
|
func (Callable): A python function that takes a Tensor or a Tensor
|
|
sequence as inputs and returns a Tensor with shape
|
|
``[batch_size, 1]`` with batch or ``[1]`` without batch.
|
|
xs (Tensor|Sequence(Tensor)): The input Tensor or Tensor sequence of
|
|
the function ``func``.
|
|
is_batched (bool): If true, the first axis is batch axis. Defaults to
|
|
False.
|
|
|
|
Returns:
|
|
|
|
Hessian (Object): A python object retains the Hessian matrix.
|
|
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> def reducer(x):
|
|
... return paddle.sum(x * x)
|
|
>>> x = paddle.rand([2, 2])
|
|
>>> h = paddle.incubate.autograd.Hessian(reducer, x)
|
|
>>> print(h[:])
|
|
Tensor(shape=[4, 4], dtype=float32, place=CPUPlace(), stop_gradient=False,
|
|
[[2., 0., 0., 0.],
|
|
[0., 2., 0., 0.],
|
|
[0., 0., 2., 0.],
|
|
[0., 0., 0., 2.]])
|
|
|
|
"""
|
|
|
|
symbolic: Jacobian
|
|
|
|
def __init__(
|
|
self,
|
|
func: Callable[..., TensorOrTensors],
|
|
xs: TensorOrTensors,
|
|
is_batched: bool = False,
|
|
) -> None:
|
|
def _jac_func(*xs):
|
|
jac = Jacobian(func, xs, is_batched=is_batched)
|
|
if (is_batched and jac.shape[1] != 1) or (
|
|
not is_batched and jac.shape[0] != 1
|
|
):
|
|
raise RuntimeError(
|
|
"The function given to Hessian should return as single element Tensor or batched single element Tensor."
|
|
)
|
|
return jac[:, 0, :] if is_batched else jac[0, :]
|
|
|
|
self.symbolic = Jacobian(_jac_func, xs, is_batched=is_batched)
|
|
|
|
def __getitem__(self, indexes: int | slice) -> Tensor:
|
|
return self.symbolic[indexes]
|
|
|
|
@property
|
|
def shape(self) -> list[int]:
|
|
"""The shape of flattened Hessian matrix."""
|
|
return self.symbolic.shape
|
|
|
|
|
|
class _Jacobian:
|
|
"""The base class for computing Jacobian matrix.
|
|
|
|
``_Jacobian`` implements the core logic of multidimensional index and lazy
|
|
evaluation for Jacobian matrix, subclass only need to overwrite following
|
|
methods:
|
|
|
|
* ``_lazy_axis()``, return the axis along which will be lazy
|
|
evaluating.
|
|
* ``_flatten(xs)``, flattens the inputs ``xs``.
|
|
* ``_evaluate(index)``, evaluates one slice along ``_lazy_axis`` .
|
|
|
|
Notes:
|
|
|
|
Because currently PaddlePaddle only support reverse differentiation by
|
|
``paddle.grad``, so lazy evaluation is only supported along the row of
|
|
Jacobian matrix, which means that slicing along row will get better
|
|
performance.
|
|
|
|
"""
|
|
|
|
def __init__(self, func, xs):
|
|
# Skip separating in prim mode temporarily, as detach and clone are not
|
|
# primitive operators.
|
|
if not framework.in_dygraph_mode() and utils.prim_enabled():
|
|
self._xs = xs
|
|
else:
|
|
self._xs = _separate(xs)
|
|
self._ys = func(*utils.as_tensors(self._xs))
|
|
self._flatten_xs = self._flatten(utils.as_tensors(self._xs))
|
|
self._flatten_ys = self._flatten(utils.as_tensors(self._ys))
|
|
self._cache = {}
|
|
|
|
@property
|
|
def shape(self):
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def _lazy_axis(self):
|
|
""" "The axis of lazily evaluated."""
|
|
raise NotImplementedError
|
|
|
|
def _lazy_indexes(self, indexes):
|
|
idx = indexes[self._lazy_axis]
|
|
return (
|
|
(idx,)
|
|
if isinstance(idx, int)
|
|
else tuple(range(idx.start, idx.stop, idx.step))
|
|
)
|
|
|
|
def _flatten(self, xs):
|
|
raise NotImplementedError
|
|
|
|
def _shifted_indexes(self, indexes, lazy_axis_size=0):
|
|
idx = indexes[self._lazy_axis]
|
|
shifted_lazy_axis_idx = (
|
|
0 if isinstance(idx, int) else slice(0, lazy_axis_size, 1)
|
|
)
|
|
return (
|
|
*indexes[: self._lazy_axis],
|
|
shifted_lazy_axis_idx,
|
|
*indexes[self._lazy_axis + 1 :],
|
|
)
|
|
|
|
def __getitem__(self, indexes):
|
|
indexes = _multi_index(indexes, self.shape)
|
|
|
|
if isinstance(indexes[self._lazy_axis], int):
|
|
other_indexes = (
|
|
indexes[: self._lazy_axis] + indexes[self._lazy_axis + 1 :]
|
|
)
|
|
return self._cached_evaluate(indexes[self._lazy_axis])[
|
|
other_indexes
|
|
]
|
|
lazy_indexes = self._lazy_indexes(indexes)
|
|
# Using concat and reshape to replace stack operator temporarily, as
|
|
# it is not a primitive operator.
|
|
shape = list(self.shape)
|
|
shape[self._lazy_axis] = len(lazy_indexes)
|
|
part_jac = paddle.concat(
|
|
[self._cached_evaluate(i) for i in lazy_indexes],
|
|
axis=self._lazy_axis,
|
|
).reshape(shape)
|
|
return part_jac[self._shifted_indexes(indexes, len(lazy_indexes))]
|
|
|
|
def _cached_evaluate(self, k):
|
|
v = self._cache.get(k)
|
|
if v is None:
|
|
v = self._evaluate(k)
|
|
self._cache[k] = v
|
|
return v
|
|
|
|
def _evaluate(self, index):
|
|
"""Evaluate one slice at along lazy axis."""
|
|
raise NotImplementedError
|
|
|
|
|
|
class _JacobianNoBatch(_Jacobian):
|
|
"""Compute Jacobian matrix without batch dimension.
|
|
Suppose the mapping is :math:`f: R^M \to R^N`, the output shape is
|
|
``(N, M)`` .
|
|
"""
|
|
|
|
def __init__(self, func, xs):
|
|
super().__init__(func, xs)
|
|
|
|
@property
|
|
def shape(self):
|
|
return (self._flatten_ys.shape[0], self._flatten_xs.shape[0])
|
|
|
|
@property
|
|
def _lazy_axis(self):
|
|
return 0
|
|
|
|
def _flatten(self, xs):
|
|
return paddle.concat(tuple(x.reshape((-1,)) for x in xs))
|
|
|
|
def _evaluate(self, row_index):
|
|
return self._flatten(
|
|
_grad(
|
|
self._flatten_ys[row_index],
|
|
self._xs,
|
|
)
|
|
)
|
|
|
|
|
|
class _JacobianBatchFirst(_Jacobian):
|
|
"""Compute Jacobian matrix with batch at first axis.
|
|
Suppose the mapping is :math:`f: R^{B,M} \to R^{B,N}`, the output shape is
|
|
``(B, N, M)`` .
|
|
"""
|
|
|
|
def __init__(self, func, xs):
|
|
super().__init__(func, xs)
|
|
|
|
@property
|
|
def shape(self):
|
|
return (
|
|
self._flatten_xs.shape[0],
|
|
self._flatten_ys.shape[1],
|
|
self._flatten_xs.shape[1],
|
|
)
|
|
|
|
@property
|
|
def _lazy_axis(self):
|
|
return 1
|
|
|
|
def _flatten(self, xs):
|
|
return paddle.concat(
|
|
tuple(x.reshape((x.shape[0], -1)) for x in utils.as_tensors(xs)), 1
|
|
)
|
|
|
|
def _evaluate(self, row_index):
|
|
return self._flatten(_grad(self._flatten_ys[:, row_index], self._xs))
|
|
|
|
|
|
def _multi_index(indexes, shape):
|
|
"""A tool for parsing N-dimensional index into a standard format.
|
|
|
|
Currently supporting following input format:
|
|
* ([positive|negative|slice], ...), the right-most elements can be
|
|
omitted.
|
|
|
|
The standard format after converted is slice tuple which contains N elements:
|
|
* ([positive|slice], ..., [positive|slice])
|
|
|
|
Notes:
|
|
Ellipsis indexes such as ``(..., i), (i, ...)`` is not supported.
|
|
|
|
Args:
|
|
indexes (tuple): The input indexes.
|
|
shape (tuple): The input shape.
|
|
|
|
Returns:
|
|
tuple: The standard format index as the above description.
|
|
"""
|
|
indexes = indexes if isinstance(indexes, typing.Sequence) else (indexes,)
|
|
if any(isinstance(i, type(Ellipsis)) for i in indexes):
|
|
raise IndexError('Ellipsis index currently is not supported.')
|
|
# Fill the right-most elements.
|
|
indexes = indexes + (slice(0, None, None),) * (len(shape) - len(indexes))
|
|
# Convert to positive index.
|
|
positive_indexes = []
|
|
for i, index in enumerate(indexes):
|
|
if isinstance(index, slice):
|
|
index = slice(
|
|
index.start or 0, index.stop or shape[i], index.step or 1
|
|
)
|
|
positive_indexes.append(
|
|
slice(
|
|
index.start + shape[i] if index.start < 0 else index.start,
|
|
index.stop + shape[i] if index.stop < 0 else index.stop,
|
|
# Negative step means index backward, no need to convert to
|
|
# positive integer.
|
|
index.step,
|
|
)
|
|
)
|
|
elif isinstance(index, int):
|
|
positive_indexes.append(index + shape[i] if index < 0 else index)
|
|
else:
|
|
raise TypeError(f'Not supported index type {index}.')
|
|
return tuple(positive_indexes)
|
|
|
|
|
|
def _replace_none_with_zero_tensor(xs, refs):
|
|
if xs is None:
|
|
xs = paddle.zeros_like(refs)
|
|
xs.stop_gradient = refs.stop_gradient
|
|
return xs
|
|
elif isinstance(xs, typing.Sequence):
|
|
return tuple(
|
|
_replace_none_with_zero_tensor(x, refs[i]) for i, x in enumerate(xs)
|
|
)
|
|
else:
|
|
return xs
|
|
|
|
|
|
def _grad(ys, xs, v=None):
|
|
"""A gradient function that can be used in dynamic graph and static graph.
|
|
|
|
The ``grad`` combines ``paddle.grad`` used in dynamic graph and
|
|
``paddle.static.gradients`` used in static graph, and do following changes:
|
|
|
|
* The ``allow_unused`` flag is removed and set defaults to true internally,
|
|
none in outputs will be replaced by zero tensor.
|
|
* The ``create_graph`` flag is removed and set defaults to true internally,
|
|
only makes sense in dynamic graph.
|
|
* When xs is a single Tensor, ``paddle.grad`` returns a list which only
|
|
contains one Tensor. It may confuse users, thus in this case we improve
|
|
to return a single Tensor in _grad interface.
|
|
|
|
Args:
|
|
ys (Tensor|Sequence[Tensor]): The output tensor or tensor sequence of
|
|
the graph to compute gradients.
|
|
xs (Tensor|Sequence[Tensor]): The input tensor or tensor sequence of the graph to
|
|
compute gradients. The returned values of this API are the
|
|
gradients of inputs .
|
|
v (Tensor|Sequence[Tensor]|None,optional): The initial gradient values
|
|
of outputs . If grad_outputs is None, the initial gradient values of
|
|
outputs would be Tensors filled with 1; if grad_outputs is not None,
|
|
it must have the same length as outputs , and in this case, the
|
|
initial gradient value of the i-th outputs would be: (1) a Tensor
|
|
filled with 1 when the i-th element of grad_outputs is None;
|
|
(2) the i-th element of grad_outputs when the i-th element of
|
|
grad_outputs is a Tensor. Default None.
|
|
|
|
Returns:
|
|
Tensor|tuple[Tensor]: Tensor or a tuple of Tensors, whose length is the
|
|
same as the Tensor number inside inputs, and the i-th returned
|
|
Tensor is the sum of gradients of outputs with respect to the i-th
|
|
inputs.
|
|
"""
|
|
if framework.in_dygraph_mode():
|
|
# paddle.grad returns a list though the inputs is a single Tensor. The
|
|
# follow code snippet fixes the problem by return the first element of
|
|
# xs_grad when the xs is a single Tensor.
|
|
xs_grad = paddle.grad(ys, xs, v, create_graph=True, allow_unused=True)
|
|
if (
|
|
isinstance(xs, (paddle.base.framework.Variable, paddle.pir.Value))
|
|
and isinstance(xs_grad, typing.Sequence)
|
|
and len(xs_grad) > 0
|
|
):
|
|
xs_grad = xs_grad[0]
|
|
else:
|
|
xs_grad = paddle.incubate.autograd.grad(ys, xs, v)
|
|
return _replace_none_with_zero_tensor(xs_grad, xs)
|
|
|
|
|
|
def _separate(xs):
|
|
"""
|
|
``_separate`` separates ``xs`` from the computation graph through ``clone``
|
|
or ``detach`` .
|
|
|
|
Internally, ``paddle.grad(xs, ys)`` is stateful API implemented based on
|
|
computational graph, which will reduce gradients along all path from ys to xs.
|
|
|
|
However, functional autograd API such as ``vjp``, ``jvp`` is stateless, and
|
|
only compute gradients with a given ``func`` .
|
|
|
|
For example, given a ``func`` :math:`y0=f(x0)`, supposing forward path is:
|
|
``x0 -> y0``, ``x0 -> x1 -> y0`` .
|
|
``paddle.grad(y0, x0)`` will reduce gradients along ``y0->x0`` and
|
|
``y0->x1->x0``, and ``vjp`` only need reduce along ``y0->x0``.
|
|
|
|
So, it's needed to clone or detach xs for breaking the dependencies with
|
|
other variables.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.incubate.autograd.functional import _separate
|
|
|
|
>>> def func(x, y):
|
|
... return x * y
|
|
>>> x = paddle.ones((1,))
|
|
>>> x.stop_gradient = False
|
|
|
|
>>> y = func(x, x)
|
|
>>> print(paddle.grad(y, x))
|
|
[Tensor(shape=[1], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[2.])]
|
|
|
|
>>> x1, x2 = _separate((x, x))
|
|
>>> y = func(x1, x2)
|
|
>>> print(paddle.grad(y, x1))
|
|
[Tensor(shape=[1], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[1.])]
|
|
|
|
"""
|
|
if isinstance(xs, typing.Sequence):
|
|
return tuple(_single_separate(x) for x in xs)
|
|
else:
|
|
return _single_separate(xs)
|
|
|
|
|
|
def _single_separate(x):
|
|
if x is None: # x maybe none because grad input's v defaults to none.
|
|
return x
|
|
if not x.stop_gradient:
|
|
return paddle.assign(x)
|
|
else: # use detach to share memory when no need gradients.
|
|
x = x.detach()
|
|
x.stop_gradient = False
|
|
return x
|
|
return x
|
|
|
|
|
|
def _check_inputs(func, xs, v=None):
|
|
if not callable(func):
|
|
raise TypeError(f"Expected 'fun' is Callable, but got {type(func)}.")
|
|
|
|
if not isinstance(
|
|
xs, (framework.Variable, typing.Sequence, paddle.pir.Value)
|
|
):
|
|
raise TypeError(
|
|
f"Expected 'xs' is a Tensor|Sequence[Tensor], but got {type(xs)}."
|
|
)
|
|
if isinstance(xs, typing.Sequence) and not all(
|
|
isinstance(x, (framework.Variable, paddle.pir.Value)) for x in xs
|
|
):
|
|
raise TypeError("All elements of 'xs' should be Tensor.")
|
|
|
|
if not isinstance(
|
|
v, (framework.Variable, typing.Sequence, type(None), paddle.pir.Value)
|
|
):
|
|
raise TypeError(
|
|
f"Expected 'v' is Tensor|Sequence[Tensor]|None, but got {type(v)}."
|
|
)
|
|
|
|
if isinstance(v, typing.Sequence) and not all(
|
|
isinstance(e, (framework.Variable, paddle.pir.Value)) for e in v
|
|
):
|
|
raise TypeError("All elements of 'xs' should be Tensor.")
|
|
|
|
|
|
def _check_v_shape(v, refs):
|
|
if v is None:
|
|
return
|
|
|
|
v, refs = utils.as_tensors(v), utils.as_tensors(refs)
|
|
if len(refs) != len(v):
|
|
raise RuntimeError(
|
|
f"The argument v is a tuple of invalid length:"
|
|
f"should be {len(refs)} but got {len(v)}."
|
|
)
|
|
|
|
for index, (element_v, element_ref) in enumerate(zip(v, refs)):
|
|
if element_v.shape != element_ref.shape:
|
|
raise RuntimeError(
|
|
f"The v[{index}] has invalid shape: should "
|
|
f"be {element_ref.shape} but got {element_v.shape}."
|
|
)
|