wehub-resource-sync
246 lines
8.5 KiB
Python
246 lines
8.5 KiB
Python
# 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
|