jax.nn.initializers.he_normal#

jax.nn.initializers.he_normal(in_axis=-2, out_axis=-1, batch_axis=(), dtype=<class 'jax.numpy.float64'>)[原始碼]#

建立 He 正規化初始器 (又稱 Kaiming 正規化初始器)。

He 正規化初始器是 jax.nn.initializers.variance_scaling() 的特化，其中 scale = 2.0、mode="fan_in" 和 distribution="truncated_normal"。

參數:

  • in_axis (int | Sequence[int]) – 權重陣列中輸入維度的軸或軸序列。

  • out_axis (int | Sequence[int]) – 權重陣列中輸出維度的軸或軸序列。

  • batch_axis (Sequence[int]) – 權重陣列中應忽略的軸或軸序列。

  • dtype (DTypeLikeInexact) – 權重的 dtype。

回傳:

一個初始器。

回傳類型:

Initializer

範例

>>> import jax, jax.numpy as jnp
>>> initializer = jax.nn.initializers.he_normal()
>>> initializer(jax.random.key(42), (2, 3), jnp.float32)  
Array([[ 0.6604483 ,  1.1900088 ,  1.2047218 ],
       [-0.87225807, -0.95321447,  0.1369438 ]], dtype=float32)