AutoDiff(自动微分) 模块

ppsci.autodiff.ad

This module is adapted from https://github.com/lululxvi/deepxde

jacobian: Callable[['paddle.Tensor', Union['paddle.Tensor', List['paddle.Tensor']], int, Optional[int], Optional[bool], bool], Union['paddle.Tensor', List['paddle.Tensor']]] = Jacobians() module-attribute

hessian: Callable[['paddle.Tensor', 'paddle.Tensor', Optional[int], int, int, Optional['paddle.Tensor'], Optional[bool], bool], 'paddle.Tensor'] = Hessians() module-attribute

Jacobians

Compute multiple Jacobians.

\[ \rm Jacobian(ys, xs, i, j) = \dfrac{\partial ys_i}{\partial xs_j} \]

A new instance will be created for a new pair of (output, input). For the (output, input) pair that has been computed before, it will reuse the previous instance, rather than creating a new one.

Source code in ppsci/autodiff/ad.py

class Jacobians:
    r"""Compute multiple Jacobians.

    $$
    \rm Jacobian(ys, xs, i, j) = \dfrac{\partial ys_i}{\partial xs_j}
    $$

    A new instance will be created for a new pair of (output, input). For the (output,
    input) pair that has been computed before, it will reuse the previous instance,
    rather than creating a new one.
    """

    def __init__(self):
        self.Js = {}

    def __call__(
        self,
        ys: "paddle.Tensor",
        xs: Union["paddle.Tensor", List["paddle.Tensor"]],
        i: int = 0,
        j: Optional[int] = None,
        retain_graph: Optional[bool] = None,
        create_graph: bool = True,
    ) -> Union["paddle.Tensor", List["paddle.Tensor"]]:
        """Compute jacobians for given ys and xs.

        Args:
            ys (paddle.Tensor): Output tensor.
            xs (Union[paddle.Tensor, List[paddle.Tensor]]): Input tensor(s).
            i (int, optional): i-th output variable. Defaults to 0.
            j (Optional[int]): j-th input variable. Defaults to None.
            retain_graph (Optional[bool]): Whether to retain the forward graph which
                is used to calculate the gradient. When it is True, the graph would
                be retained, in which way users can calculate backward twice for the
                same graph. When it is False, the graph would be freed. Default None,
                which means it is equal to `create_graph`.
            create_graph (bool, optional): Whether to create the gradient graphs of
                the computing process. When it is True, higher order derivatives are
                supported to compute; when it is False, the gradient graphs of the
                computing process would be discarded. Default False.

        Returns:
            paddle.Tensor: Jacobian matrix of ys[i] to xs[j].

        Examples:
            >>> import paddle
            >>> import ppsci
            >>> x = paddle.randn([4, 1])
            >>> x.stop_gradient = False
            >>> y = x * x
            >>> dy_dx = ppsci.autodiff.jacobian(y, x)
            >>> print(dy_dx.shape)
            [4, 1]
        """
        if not isinstance(xs, (list, tuple)):
            key = (ys, xs)
            if key not in self.Js:
                self.Js[key] = _Jacobian(ys, xs)
            return self.Js[key](i, j, retain_graph, create_graph)
        else:
            xs_require = [xs[i] for i in range(len(xs)) if (ys, xs[i]) not in self.Js]
            grads_require = paddle.grad(
                ys,
                xs_require,
                create_graph=create_graph,
                retain_graph=retain_graph,
            )

            idx = 0
            Js_list = []
            for k, xs_ in enumerate(xs):
                key = (ys, xs_)
                assert xs_.shape[-1] == 1, (
                    f"The last dim of each xs should be 1, but xs[{k}] has shape "
                    f"{xs_.shape}"
                )
                if key not in self.Js:
                    self.Js[key] = _Jacobian(ys, xs_, {0: grads_require[idx]})
                    idx += 1
                Js_list.append(self.Js[key](i, j, retain_graph, create_graph))
            return Js_list

    def _clear(self):
        """Clear cached Jacobians."""
        self.Js = {}

__call__(ys, xs, i=0, j=None, retain_graph=None, create_graph=True)

Compute jacobians for given ys and xs.

Parameters:

Name	Type	Description	Default
ys	Tensor	Output tensor.	required
xs	Union[Tensor, List[Tensor]]	Input tensor(s).	required
i	int	i-th output variable. Defaults to 0.	0
j	Optional[int]	j-th input variable. Defaults to None.	None
retain_graph	Optional[bool]	Whether to retain the forward graph which is used to calculate the gradient. When it is True, the graph would be retained, in which way users can calculate backward twice for the same graph. When it is False, the graph would be freed. Default None, which means it is equal to create_graph.	None
create_graph	bool	Whether to create the gradient graphs of the computing process. When it is True, higher order derivatives are supported to compute; when it is False, the gradient graphs of the computing process would be discarded. Default False.	True

Returns:

Type	Description
Union['paddle.Tensor', List['paddle.Tensor']]	paddle.Tensor: Jacobian matrix of ys[i] to xs[j].

Examples:

>>> import paddle
>>> import ppsci
>>> x = paddle.randn([4, 1])
>>> x.stop_gradient = False
>>> y = x * x
>>> dy_dx = ppsci.autodiff.jacobian(y, x)
>>> print(dy_dx.shape)
[4, 1]

Source code in ppsci/autodiff/ad.py

def __call__(
    self,
    ys: "paddle.Tensor",
    xs: Union["paddle.Tensor", List["paddle.Tensor"]],
    i: int = 0,
    j: Optional[int] = None,
    retain_graph: Optional[bool] = None,
    create_graph: bool = True,
) -> Union["paddle.Tensor", List["paddle.Tensor"]]:
    """Compute jacobians for given ys and xs.

    Args:
        ys (paddle.Tensor): Output tensor.
        xs (Union[paddle.Tensor, List[paddle.Tensor]]): Input tensor(s).
        i (int, optional): i-th output variable. Defaults to 0.
        j (Optional[int]): j-th input variable. Defaults to None.
        retain_graph (Optional[bool]): Whether to retain the forward graph which
            is used to calculate the gradient. When it is True, the graph would
            be retained, in which way users can calculate backward twice for the
            same graph. When it is False, the graph would be freed. Default None,
            which means it is equal to `create_graph`.
        create_graph (bool, optional): Whether to create the gradient graphs of
            the computing process. When it is True, higher order derivatives are
            supported to compute; when it is False, the gradient graphs of the
            computing process would be discarded. Default False.

    Returns:
        paddle.Tensor: Jacobian matrix of ys[i] to xs[j].

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> x = paddle.randn([4, 1])
        >>> x.stop_gradient = False
        >>> y = x * x
        >>> dy_dx = ppsci.autodiff.jacobian(y, x)
        >>> print(dy_dx.shape)
        [4, 1]
    """
    if not isinstance(xs, (list, tuple)):
        key = (ys, xs)
        if key not in self.Js:
            self.Js[key] = _Jacobian(ys, xs)
        return self.Js[key](i, j, retain_graph, create_graph)
    else:
        xs_require = [xs[i] for i in range(len(xs)) if (ys, xs[i]) not in self.Js]
        grads_require = paddle.grad(
            ys,
            xs_require,
            create_graph=create_graph,
            retain_graph=retain_graph,
        )

        idx = 0
        Js_list = []
        for k, xs_ in enumerate(xs):
            key = (ys, xs_)
            assert xs_.shape[-1] == 1, (
                f"The last dim of each xs should be 1, but xs[{k}] has shape "
                f"{xs_.shape}"
            )
            if key not in self.Js:
                self.Js[key] = _Jacobian(ys, xs_, {0: grads_require[idx]})
                idx += 1
            Js_list.append(self.Js[key](i, j, retain_graph, create_graph))
        return Js_list

Hessians

Compute multiple Hessians.

\[ \rm Hessian(ys, xs, component, i, j) = \dfrac{\partial ys_{component}}{\partial xs_i \partial xs_j} \]

A new instance will be created for a new pair of (output, input). For the (output, input) pair that has been computed before, it will reuse the previous instance, rather than creating a new one.

Source code in ppsci/autodiff/ad.py

class Hessians:
    r"""Compute multiple Hessians.

    $$
    \rm Hessian(ys, xs, component, i, j) = \dfrac{\partial ys_{component}}{\partial xs_i \partial xs_j}
    $$

    A new instance will be created for a new pair of (output, input). For the (output,
    input) pair that has been computed before, it will reuse the previous instance,
    rather than creating a new one.
    """

    def __init__(self):
        self.Hs = {}

    def __call__(
        self,
        ys: "paddle.Tensor",
        xs: "paddle.Tensor",
        component: Optional[int] = None,
        i: int = 0,
        j: int = 0,
        grad_y: Optional["paddle.Tensor"] = None,
        retain_graph: Optional[bool] = None,
        create_graph: bool = True,
    ) -> "paddle.Tensor":
        """Compute hessian matrix for given ys and xs.

        Args:
            ys (paddle.Tensor): Output tensor.
            xs (paddle.Tensor): Input tensor.
            component (Optional[int]): If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]`
                is used to compute the Hessian. Do not use if `y` has the shape (batch_size,
                1). Defaults to None.
            i (int, optional): I-th input variable. Defaults to 0.
            j (int, optional): J-th input variable. Defaults to 0.
            grad_y (Optional[paddle.Tensor]): The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid
                duplicate computation. Defaults to None.
            retain_graph (Optional[bool]): Whether to retain the forward graph which
                is used to calculate the gradient. When it is True, the graph would
                be retained, in which way users can calculate backward twice for the
                same graph. When it is False, the graph would be freed. Default None,
                which means it is equal to `create_graph`.
            create_graph (bool, optional): Whether to create the gradient graphs of
                the computing process. When it is True, higher order derivatives are
                supported to compute; when it is False, the gradient graphs of the
                computing process would be discarded. Default False.

        Returns:
            paddle.Tensor: Hessian matrix.

        Examples:
            >>> import paddle
            >>> import ppsci
            >>> x = paddle.randn([4, 3])
            >>> x.stop_gradient = False
            >>> y = (x * x).sin()
            >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0)
            >>> print(dy_dxx.shape)
            [4, 1]
        """
        key = (ys, xs, component)
        if key not in self.Hs:
            self.Hs[key] = _Hessian(ys, xs, component=component, grad_y=grad_y)
        return self.Hs[key](i, j, retain_graph, create_graph)

    def _clear(self):
        """Clear cached Hessians."""
        self.Hs = {}

__call__(ys, xs, component=None, i=0, j=0, grad_y=None, retain_graph=None, create_graph=True)

Compute hessian matrix for given ys and xs.

Parameters:

Name	Type	Description	Default
ys	Tensor	Output tensor.	required
xs	Tensor	Input tensor.	required
component	Optional[int]	If y has the shape (batch_size, dim_y > 1), then y[:, component] is used to compute the Hessian. Do not use if y has the shape (batch_size, 1). Defaults to None.	None
i	int	I-th input variable. Defaults to 0.	0
j	int	J-th input variable. Defaults to 0.	0
grad_y	Optional[Tensor]	The gradient of y w.r.t. xs. Provide grad_y if known to avoid duplicate computation. Defaults to None.	None
retain_graph	Optional[bool]	Whether to retain the forward graph which is used to calculate the gradient. When it is True, the graph would be retained, in which way users can calculate backward twice for the same graph. When it is False, the graph would be freed. Default None, which means it is equal to create_graph.	None
create_graph	bool	Whether to create the gradient graphs of the computing process. When it is True, higher order derivatives are supported to compute; when it is False, the gradient graphs of the computing process would be discarded. Default False.	True

Returns:

Type	Description
'paddle.Tensor'	paddle.Tensor: Hessian matrix.

Examples:

>>> import paddle
>>> import ppsci
>>> x = paddle.randn([4, 3])
>>> x.stop_gradient = False
>>> y = (x * x).sin()
>>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0)
>>> print(dy_dxx.shape)
[4, 1]

Source code in ppsci/autodiff/ad.py

def __call__(
    self,
    ys: "paddle.Tensor",
    xs: "paddle.Tensor",
    component: Optional[int] = None,
    i: int = 0,
    j: int = 0,
    grad_y: Optional["paddle.Tensor"] = None,
    retain_graph: Optional[bool] = None,
    create_graph: bool = True,
) -> "paddle.Tensor":
    """Compute hessian matrix for given ys and xs.

    Args:
        ys (paddle.Tensor): Output tensor.
        xs (paddle.Tensor): Input tensor.
        component (Optional[int]): If `y` has the shape (batch_size, dim_y > 1), then `y[:, component]`
            is used to compute the Hessian. Do not use if `y` has the shape (batch_size,
            1). Defaults to None.
        i (int, optional): I-th input variable. Defaults to 0.
        j (int, optional): J-th input variable. Defaults to 0.
        grad_y (Optional[paddle.Tensor]): The gradient of `y` w.r.t. `xs`. Provide `grad_y` if known to avoid
            duplicate computation. Defaults to None.
        retain_graph (Optional[bool]): Whether to retain the forward graph which
            is used to calculate the gradient. When it is True, the graph would
            be retained, in which way users can calculate backward twice for the
            same graph. When it is False, the graph would be freed. Default None,
            which means it is equal to `create_graph`.
        create_graph (bool, optional): Whether to create the gradient graphs of
            the computing process. When it is True, higher order derivatives are
            supported to compute; when it is False, the gradient graphs of the
            computing process would be discarded. Default False.

    Returns:
        paddle.Tensor: Hessian matrix.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> x = paddle.randn([4, 3])
        >>> x.stop_gradient = False
        >>> y = (x * x).sin()
        >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0)
        >>> print(dy_dxx.shape)
        [4, 1]
    """
    key = (ys, xs, component)
    if key not in self.Hs:
        self.Hs[key] = _Hessian(ys, xs, component=component, grad_y=grad_y)
    return self.Hs[key](i, j, retain_graph, create_graph)

clear()

Clear cached Jacobians and Hessians.

Examples:

>>> import paddle
>>> import ppsci
>>> x = paddle.randn([4, 3])
>>> x.stop_gradient = False
>>> y = (x * x).sin()
>>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0)
>>> ppsci.autodiff.clear()
>>> print(ppsci.autodiff.hessian.Hs)
{}

Source code in ppsci/autodiff/ad.py

def clear():
    """Clear cached Jacobians and Hessians.

    Examples:
        >>> import paddle
        >>> import ppsci
        >>> x = paddle.randn([4, 3])
        >>> x.stop_gradient = False
        >>> y = (x * x).sin()
        >>> dy_dxx = ppsci.autodiff.hessian(y, x, component=0)
        >>> ppsci.autodiff.clear()
        >>> print(ppsci.autodiff.hessian.Hs)
        {}
    """
    jacobian._clear()
    hessian._clear()