jax.nn.initializers.he_uniform

目录

jax.nn.initializers.he_uniform#

jax.nn.initializers.he_uniform(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)[源代码]#

构建一个He均匀初始化器（又名Kaiming均匀初始化器）。

He uniform initializer 是 jax.nn.initializers.variance_scaling() 的一个特化，其中 scale = 2.0，mode="fan_in"，和 distribution="uniform"。

参数:

  • in_axis (int | Sequence[int]) – 权重数组中输入维度的轴或轴序列。

  • out_axis (int | Sequence[int]) – 权重数组中输出维度的轴或轴序列。

  • batch_axis (Sequence[int]) – 权重数组中应忽略的轴或轴序列。

  • dtype (DTypeLikeInexact) – 权重的数据类型。

返回:

一个初始化器。

返回类型:

Initializer

示例：

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_uniform()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.79611576,  1.2789248 ,  1.2896855 ],
       [-1.0108745 , -1.0855657 ,  0.17398663]], dtype=float32)
目录