kaiming_uniform

shrew::nn::init

Function kaiming_uniform 

Source 
pub fn kaiming_uniform<B>(
    shape: impl Into<Shape>,
    a: f64,
    mode: FanMode,
    dtype: DType,
    device: &<B as Backend>::Device,
) -> Result<Tensor<B>, Error>
where
    B: Backend,
Expand description

Kaiming (He) uniform initialization.

Draws from U(-bound, bound) where bound = sqrt(3 * gain² / fan). Designed for layers followed by ReLU (or variants).

§Arguments

  • a: negative slope of the rectifier (0 for ReLU, 0.01 for LeakyReLU)
  • mode: FanIn or FanOut