473 lines
13 KiB
Python
473 lines
13 KiB
Python
# Copyright (c) 2025 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
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
from typing import TypeVar
|
|
|
|
from typing_extensions import ParamSpec
|
|
|
|
_InputT = ParamSpec("_InputT")
|
|
_RetT = TypeVar("_RetT")
|
|
|
|
import math
|
|
import warnings
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
|
|
from ..base.framework import in_dygraph_mode, in_pir_mode
|
|
from .initializer.constant import Constant
|
|
from .initializer.dirac import Dirac
|
|
from .initializer.initializer import calculate_gain # noqa: F401
|
|
from .initializer.kaiming import KaimingNormal, KaimingUniform
|
|
from .initializer.normal import Normal, TruncatedNormal
|
|
from .initializer.orthogonal import Orthogonal
|
|
from .initializer.uniform import Uniform
|
|
from .initializer.xavier import XavierNormal, XavierUniform
|
|
|
|
|
|
def _calculate_fan_in_and_fan_out(var: paddle.Tensor) -> tuple[int, int]:
|
|
"""Compute the fan_in and the fan_out for layers
|
|
|
|
This method computes the fan_in and the fan_out
|
|
for neural network layers, if not specified. It is
|
|
not possible to perfectly estimate fan_in and fan_out.
|
|
This method will estimate it correctly for matrix multiply and
|
|
convolutions.
|
|
|
|
Args:
|
|
var: variable for which fan_in and fan_out have to be computed.
|
|
|
|
Returns:
|
|
tuple of two integers (fan_in, fan_out).
|
|
"""
|
|
shape = var.shape
|
|
if not shape or len(shape) == 0:
|
|
fan_in = fan_out = 1
|
|
elif len(shape) == 1:
|
|
fan_in = fan_out = shape[0]
|
|
elif len(shape) == 2:
|
|
# This is the case for simple matrix multiply
|
|
fan_in = shape[0]
|
|
fan_out = shape[1]
|
|
else:
|
|
# Assume this to be a convolutional kernel
|
|
# In PaddlePaddle, the shape of the kernel is like:
|
|
# [num_filters, num_filter_channels, ...] where the remaining
|
|
# dimensions are the filter_size
|
|
receptive_field_size = np.prod(shape[2:])
|
|
fan_in = int(shape[1] * receptive_field_size)
|
|
fan_out = int(shape[0] * receptive_field_size)
|
|
return (fan_in, fan_out)
|
|
|
|
|
|
def kaiming_uniform_(
|
|
tensor: paddle.Tensor,
|
|
a: float = 0,
|
|
mode: str = "fan_in",
|
|
nonlinearity: str = "leaky_relu",
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using Kaiming uniform method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
a (float, optional): The negative slope of the rectifier used after this layer.
|
|
Defaults to 0.
|
|
mode (str, optional): Mode to compute the fan. Choose from ["fan_in", "fan_out"].
|
|
When set to 'fan_in', the fan_in parameter is used for initialization.
|
|
When set to 'fan_out', the out_features of trainable Tensor will be used.
|
|
Default is 'fan_in'.
|
|
nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu".
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = KaimingUniform(
|
|
negative_slope=a, nonlinearity=nonlinearity, mode=mode
|
|
)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def kaiming_normal_(
|
|
tensor: paddle.Tensor,
|
|
a: float = 0,
|
|
mode: str = "fan_in",
|
|
nonlinearity: str = "leaky_relu",
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using Kaiming normal method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
a (float, optional): The negative slope of the rectifier used after this layer.
|
|
Defaults to 0.
|
|
mode (str, optional): Mode to compute the fan. Choose from ["fan_in", "fan_out"].
|
|
When set to 'fan_in', the fan_in parameter is used for initialization.
|
|
When set to 'fan_out', the out_features of trainable Tensor will be used.
|
|
Default is 'fan_in'.
|
|
nonlinearity (str, optional): Nonlinearity method name. Defaults to "leaky_relu".
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = KaimingNormal(negative_slope=a, nonlinearity=nonlinearity, mode=mode)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def xavier_uniform_(
|
|
tensor: paddle.Tensor,
|
|
gain: float = 1.0,
|
|
fan_in: float | None = None,
|
|
fan_out: float | None = None,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using Xavier uniform method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
gain (float, optional): Scaling Tensor. Default is 1.0.
|
|
fan_in (float|None, optional): fan_in for Xavier initialization, which is
|
|
inferred from the Tensor. Default is None.
|
|
fan_out (float|None, optional): fan_out for Xavier initialization, which is
|
|
inferred from the Tensor. Default is None.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = XavierUniform(
|
|
gain=gain,
|
|
fan_in=fan_in,
|
|
fan_out=fan_out,
|
|
)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def xavier_normal_(
|
|
tensor: paddle.Tensor,
|
|
gain: float = 1.0,
|
|
fan_in: float | None = None,
|
|
fan_out: float | None = None,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using Xavier normal method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
gain (float, optional): Scaling Tensor. Default is 1.0.
|
|
fan_in (float|None, optional): fan_in for Xavier initialization, which is
|
|
inferred from the Tensor. Default is None.
|
|
fan_out (float|None, optional): fan_out for Xavier initialization, which is
|
|
inferred from the Tensor. Default is None.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = XavierNormal(
|
|
gain=gain,
|
|
fan_in=fan_in,
|
|
fan_out=fan_out,
|
|
)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def uniform_(
|
|
tensor: paddle.Tensor,
|
|
a: float = 0.0,
|
|
b: float = 1.0,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using uniform method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
low (float, optional): Lower boundary of the uniform distribution. Default is :math:`-1.0`.
|
|
high (float, optional): Upper boundary of the uniform distribution. Default is :math:`1.0`.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Uniform(low=a, high=b)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def normal_(
|
|
tensor: paddle.Tensor,
|
|
mean: float = 0.0,
|
|
std: float = 1.0,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using normal method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
mean (float|complex, optional): mean of the normal distribution. Default is 0.0.
|
|
std (float, optional): standard deviation of the normal distribution. Default is 1.0.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Normal(mean=mean, std=std)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def trunc_normal_(
|
|
tensor: paddle.Tensor,
|
|
mean: float = 0.0,
|
|
std: float = 1.0,
|
|
a: float = -2.0,
|
|
b: float = 2.0,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using truncated normal method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
mean (float|complex, optional): mean of the normal distribution. Default is 0.0.
|
|
std (float, optional): standard deviation of the normal distribution. Default is 1.0.
|
|
a (float, optional): The minimum cutoff value. Default is -2.0.
|
|
b (float, optional): The maximum cutoff value. Default is 2.0.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = TruncatedNormal(mean=mean, std=std, a=a, b=b)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def constant_(
|
|
tensor: paddle.Tensor,
|
|
val: float,
|
|
) -> paddle.Tensor:
|
|
"""Modify tensor inplace using constant method.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
value (float32|float64, optional): constant value to initialize the parameter.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Constant(value=val)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def ones_(
|
|
tensor: paddle.Tensor,
|
|
) -> paddle.Tensor:
|
|
"""Fill the input Tensor with the scalar value 1.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Constant(value=1.0)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def zeros_(
|
|
tensor: paddle.Tensor,
|
|
) -> paddle.Tensor:
|
|
"""Fill the input Tensor with the scalar value 0.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Constant(value=0.0)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def dirac_(
|
|
tensor: paddle.Tensor,
|
|
groups: int = 1,
|
|
) -> paddle.Tensor:
|
|
"""Initialize the 3D/4D/5D Tensor with Dirac delta function.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
groups (int|None, optional): 0-dimension of the Tensor will be divided by groups,
|
|
each group has the same value. Default: 1.
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Dirac(groups=groups)
|
|
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def eye_(
|
|
tensor: paddle.Tensor,
|
|
) -> paddle.Tensor:
|
|
"""Fill the 2-dimensional input Tensor with the identity matrix.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
|
|
if len(tensor.shape) != 2:
|
|
raise AssertionError(
|
|
f"Only support 2 dimensional tensor, but got {len(tensor.shape)}."
|
|
)
|
|
|
|
if in_dygraph_mode():
|
|
new_tensor = paddle.eye(
|
|
tensor.shape[0], tensor.shape[1], dtype=tensor.dtype
|
|
)
|
|
new_tensor._share_underline_tensor_to(tensor)
|
|
return tensor
|
|
elif in_pir_mode():
|
|
new_tensor = paddle.eye(
|
|
tensor.shape[0], tensor.shape[1], dtype=tensor.dtype
|
|
)
|
|
return new_tensor
|
|
else:
|
|
raise NotImplementedError(
|
|
'Only support run in dygraph mode or PIR mode.'
|
|
)
|
|
|
|
|
|
def orthogonal_(
|
|
tensor: paddle.Tensor,
|
|
gain: float = 1,
|
|
) -> paddle.Tensor:
|
|
"""Fill the input Tensor with a (semi) orthogonal matrix.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor.
|
|
gain(float, optional): The multiplication coefficient for initialized tensor. Default: 1.0.
|
|
Returns:
|
|
Tensor: Initialized tensor.
|
|
"""
|
|
init = Orthogonal(gain=gain)
|
|
if in_dygraph_mode():
|
|
init(tensor)
|
|
return tensor
|
|
return init(tensor)
|
|
|
|
|
|
def sparse_(
|
|
tensor: paddle.Tensor, sparsity: float, std: float = 0.01
|
|
) -> paddle.Tensor:
|
|
"""Fill the 2D input Tensor as a sparse matrix.
|
|
|
|
The non-zero elements will be drawn from the normal distribution.
|
|
|
|
Args:
|
|
tensor (Tensor): Paddle Tensor with 2 dimensions.
|
|
sparsity (float): The fraction of elements in each column to be set to zero.
|
|
std (float): the standard deviation of the normal distribution used to generate
|
|
the non-zero values. Default is 0.01.
|
|
|
|
Examples:
|
|
>>> tensor = paddle.empty(3, 5)
|
|
>>> result = paddle.nn.init.sparse_(tensor, sparsity=0.1)
|
|
"""
|
|
if tensor.ndimension() != 2:
|
|
raise ValueError("Only tensors with 2 dimensions are supported")
|
|
rows, cols = tensor.shape
|
|
num_zeros = math.ceil(sparsity * rows)
|
|
|
|
with paddle.no_grad():
|
|
tensor = normal_(tensor, mean=0, std=std)
|
|
for col_idx in range(cols):
|
|
row_indices = paddle.randperm(rows)
|
|
zero_indices = row_indices[:num_zeros]
|
|
tensor[zero_indices, col_idx] = 0
|
|
return tensor
|
|
|
|
|
|
def _make_deprecate(func: Callable[_InputT, _RetT]) -> Callable[_InputT, _RetT]:
|
|
new_name = func.__name__
|
|
old_name = new_name[:-1]
|
|
|
|
def deprecated_init(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
|
|
warnings.warn(
|
|
f"`nn.init.{old_name}` is now deprecated in favor of `nn.init.{new_name}`.",
|
|
FutureWarning,
|
|
stacklevel=2,
|
|
)
|
|
return func(*args, **kwargs)
|
|
|
|
deprecated_init.__doc__ = rf"""
|
|
{old_name}(...)
|
|
|
|
.. warning::
|
|
This method is now deprecated in favor of :func:`paddle.nn.init.{new_name}`.
|
|
|
|
See :func:`~paddle.nn.init.{new_name}` for details."""
|
|
deprecated_init.__name__ = old_name
|
|
return deprecated_init
|
|
|
|
|
|
uniform = _make_deprecate(uniform_)
|
|
normal = _make_deprecate(normal_)
|
|
constant = _make_deprecate(constant_)
|
|
eye = _make_deprecate(eye_)
|
|
dirac = _make_deprecate(dirac_)
|
|
xavier_uniform = _make_deprecate(xavier_uniform_)
|
|
xavier_normal = _make_deprecate(xavier_normal_)
|
|
kaiming_uniform = _make_deprecate(kaiming_uniform_)
|
|
kaiming_normal = _make_deprecate(kaiming_normal_)
|
|
orthogonal = _make_deprecate(orthogonal_)
|
|
sparse = _make_deprecate(sparse_)
|