chore: import upstream snapshot with attribution

python/paddle/nn/utils/spectral_norm_hook.py

@@ -0,0 +1,245 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+
+from typing import TYPE_CHECKING
+
+import paddle
+
+from .. import functional as F
+from ..layer.common import Linear
+from ..layer.conv import Conv1DTranspose, Conv2DTranspose, Conv3DTranspose
+
+if TYPE_CHECKING:
+    from typing_extensions import Never
+
+    from paddle import Tensor
+    from paddle.nn import Layer
+
+__all__ = []
+
+
+def normal_(x: Tensor, mean: float = 0.0, std: float = 1.0) -> Tensor:
+    temp_value = paddle.normal(mean, std, shape=x.shape)
+    paddle.assign(temp_value, x)
+    return x
+
+
+class SpectralNorm:
+    name: str
+    dim: int
+    n_power_iterations: int
+    eps: float
+
+    def __init__(
+        self,
+        name: str = 'weight',
+        n_power_iterations: int = 1,
+        dim: int = 0,
+        eps: float = 1e-12,
+    ) -> None:
+        self.name = name
+        self.dim = dim
+        if n_power_iterations <= 0:
+            raise ValueError(
+                'Expected n_power_iterations to be positive, but '
+                f'got n_power_iterations={n_power_iterations}'
+            )
+        self.n_power_iterations = n_power_iterations
+        self.eps = eps
+
+    def reshape_weight_to_matrix(self, weight: Tensor) -> Tensor:
+        weight_mat = weight
+        if self.dim != 0:
+            # transpose dim to front
+            weight_mat = weight_mat.transpose(
+                [self.dim]
+                + [d for d in range(weight_mat.dim()) if d != self.dim]
+            )
+
+        height = weight_mat.shape[0]
+
+        return weight_mat.reshape([height, -1])
+
+    def compute_weight(self, layer: Layer, do_power_iteration: bool) -> Tensor:
+        weight = getattr(layer, self.name + '_orig')
+        u = getattr(layer, self.name + '_u')
+        v = getattr(layer, self.name + '_v')
+        weight_mat = self.reshape_weight_to_matrix(weight)
+
+        if do_power_iteration:
+            with paddle.no_grad():
+                for _ in range(self.n_power_iterations):
+                    paddle.assign(
+                        F.normalize(
+                            paddle.matmul(
+                                weight_mat,
+                                u,
+                                transpose_x=True,
+                                transpose_y=False,
+                            ),
+                            axis=0,
+                            epsilon=self.eps,
+                        ),
+                        v,
+                    )
+
+                    paddle.assign(
+                        F.normalize(
+                            paddle.matmul(weight_mat, v),
+                            axis=0,
+                            epsilon=self.eps,
+                        ),
+                        u,
+                    )
+                if self.n_power_iterations > 0:
+                    u = u.clone()
+                    v = v.clone()
+
+        sigma = paddle.dot(u, paddle.mv(weight_mat, v))
+        weight = weight / sigma
+        return weight
+
+    def __call__(self, layer: Layer, inputs: Never) -> None:
+        setattr(
+            layer,
+            self.name,
+            self.compute_weight(layer, do_power_iteration=layer.training),
+        )
+
+    @staticmethod
+    def apply(
+        layer: Layer, name: str, n_power_iterations: int, dim: int, eps: float
+    ) -> SpectralNorm:
+        for k, hook in layer._forward_pre_hooks.items():
+            if isinstance(hook, SpectralNorm) and hook.name == name:
+                raise RuntimeError(
+                    "Cannot register two spectral_norm hooks on "
+                    f"the same parameter {name}"
+                )
+
+        fn = SpectralNorm(name, n_power_iterations, dim, eps)
+        weight = layer._parameters[name]
+
+        with paddle.no_grad():
+            weight_mat = fn.reshape_weight_to_matrix(weight)
+            h, w = weight_mat.shape
+
+            # randomly initialize u and v
+            u = layer.create_parameter([h])
+            u = normal_(u, 0.0, 1.0)
+            v = layer.create_parameter([w])
+            v = normal_(v, 0.0, 1.0)
+            u = F.normalize(u, axis=0, epsilon=fn.eps)
+            v = F.normalize(v, axis=0, epsilon=fn.eps)
+
+        # delete fn.name form parameters, otherwise you can not set attribute
+        del layer._parameters[fn.name]
+        layer.add_parameter(fn.name + "_orig", weight)
+        # still need to assign weight back as fn.name because all sorts of
+        # things may assume that it exists, e.g., when initializing weights.
+        # However, we can't directly assign as it could be an Parameter and
+        # gets added as a parameter. Instead, we register weight * 1.0 as a plain
+        # attribute.
+        setattr(layer, fn.name, weight * 1.0)
+        layer.register_buffer(fn.name + "_u", u)
+        layer.register_buffer(fn.name + "_v", v)
+        layer.register_forward_pre_hook(fn)
+        return fn
+
+
+def spectral_norm(
+    layer: Layer,
+    name: str = 'weight',
+    n_power_iterations: int = 1,
+    eps: float = 1e-12,
+    dim: int | None = None,
+) -> Layer:
+    r"""
+    Applies spectral normalization to a parameter according to the
+    following Calculation:
+
+    Step 1:
+    Generate vector U in shape of [H], and V in shape of [W].
+    While H is the :attr:`dim` th dimension of the input weights,
+    and W is the product result of remaining dimensions.
+
+    Step 2:
+    :attr:`n_power_iterations` should be a positive integer, do following
+    calculations with U and V for :attr:`power_iters` rounds.
+
+    .. math::
+
+        \mathbf{v} := \frac{\mathbf{W}^{T} \mathbf{u}}{\|\mathbf{W}^{T} \mathbf{u}\|_2}
+
+        \mathbf{u} := \frac{\mathbf{W} \mathbf{v}}{\|\mathbf{W} \mathbf{v}\|_2}
+
+    Step 3:
+    Calculate :math:`\sigma(\mathbf{W})` and normalize weight values.
+
+    .. math::
+
+        \sigma(\mathbf{W}) = \mathbf{u}^{T} \mathbf{W} \mathbf{v}
+
+        \mathbf{W} = \frac{\mathbf{W}}{\sigma(\mathbf{W})}
+
+
+    Refer to `Spectral Normalization <https://arxiv.org/abs/1802.05957>`_ .
+
+    Parameters:
+        layer(Layer): Layer of paddle, which has weight.
+        name(str, optional): Name of the weight parameter. Default: 'weight'.
+        n_power_iterations(int, optional): The number of power iterations to calculate spectral norm. Default: 1.
+        eps(float, optional): The epsilon for numerical stability in calculating norms. Default: 1e-12.
+        dim(int|None, optional): The index of dimension which should be permuted to the first before reshaping Input(Weight) to matrix, it should be set as 0 if Input(Weight) is the weight of fc layer, and should be set as 1 if Input(Weight) is the weight of conv layer. Default: None.
+
+    Returns:
+        Layer, the original layer with the spectral norm hook.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> from paddle.nn import Conv2D
+            >>> from paddle.nn.utils import spectral_norm
+            >>> paddle.seed(2023)
+            >>> conv = Conv2D(3, 1, 3)
+            >>> sn_conv = spectral_norm(conv)
+            >>> print(sn_conv)
+            Conv2D(3, 1, kernel_size=[3, 3], data_format=NCHW)
+            >>> # Conv2D(3, 1, kernel_size=[3, 3], data_format=NCHW)
+            >>> print(sn_conv.weight)
+            Tensor(shape=[1, 3, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=False,
+            [[[[ 0.01668976,  0.30305523,  0.11405435],
+               [-0.06765547, -0.50396705, -0.40925547],
+               [ 0.47344422,  0.03628403,  0.45277366]],
+              [[-0.15177251, -0.16305730, -0.15723954],
+               [-0.28081197, -0.09183260, -0.08081978],
+               [-0.40895155,  0.18298769, -0.29325116]],
+              [[ 0.21819633, -0.01822380, -0.50351536],
+               [-0.06262003,  0.17713565,  0.20517939],
+               [ 0.16659889, -0.14333329,  0.05228264]]]])
+
+    """
+
+    if dim is None:
+        if isinstance(
+            layer, (Conv1DTranspose, Conv2DTranspose, Conv3DTranspose, Linear)
+        ):
+            dim = 1
+        else:
+            dim = 0
+    SpectralNorm.apply(layer, name, n_power_iterations, dim, eps)
+    return layer