100 lines
2.7 KiB
Rust
100 lines
2.7 KiB
Rust
//! Parameter initialization functions
|
|
|
|
use crate::Result;
|
|
use rand::prelude::*;
|
|
use rand_distr::{Distribution, Normal, Uniform};
|
|
use rtx_tensor::Tensor;
|
|
|
|
/// Non-linearity types for initialization
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum NonLinearity {
|
|
/// Linear activation
|
|
Linear,
|
|
/// ReLU activation
|
|
ReLU,
|
|
}
|
|
|
|
impl NonLinearity {
|
|
/// Get the gain factor for this non-linearity
|
|
pub fn gain(self) -> f32 {
|
|
match self {
|
|
Self::Linear => 1.0,
|
|
Self::ReLU => (2.0_f32).sqrt(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Initialize tensor with uniform distribution
|
|
pub fn uniform(tensor: &mut Tensor, low: f32, high: f32) -> Result<()> {
|
|
let shape = tensor.shape().dims().to_vec();
|
|
let device = tensor.device().clone();
|
|
let numel = tensor.numel();
|
|
|
|
let mut rng = thread_rng();
|
|
let dist = Uniform::new(low, high);
|
|
let data: Vec<f32> = (0..numel).map(|_| dist.sample(&mut rng)).collect();
|
|
|
|
*tensor = Tensor::from_data(data, shape, &device)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize tensor with normal distribution
|
|
pub fn normal(tensor: &mut Tensor, mean: f32, std: f32) -> Result<()> {
|
|
let shape = tensor.shape().dims().to_vec();
|
|
let device = tensor.device().clone();
|
|
let numel = tensor.numel();
|
|
|
|
let mut rng = thread_rng();
|
|
let dist = Normal::new(mean, std).map_err(|e| crate::Error::InvalidParameter(e.to_string()))?;
|
|
let data: Vec<f32> = (0..numel).map(|_| dist.sample(&mut rng)).collect();
|
|
|
|
*tensor = Tensor::from_data(data, shape, &device)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize tensor with zeros
|
|
pub fn zeros(tensor: &mut Tensor) -> Result<()> {
|
|
*tensor = Tensor::zeros_like(tensor)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize tensor with ones
|
|
pub fn ones(tensor: &mut Tensor) -> Result<()> {
|
|
*tensor = Tensor::ones_like(tensor)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Kaiming uniform initialization
|
|
pub fn kaiming_uniform(
|
|
tensor: &mut Tensor,
|
|
fan_in: usize,
|
|
nonlinearity: NonLinearity,
|
|
) -> Result<()> {
|
|
let gain = nonlinearity.gain();
|
|
let bound = gain * (3.0 / fan_in as f32).sqrt();
|
|
uniform(tensor, -bound, bound)
|
|
}
|
|
|
|
/// Kaiming normal initialization
|
|
pub fn kaiming_normal(
|
|
tensor: &mut Tensor,
|
|
fan_in: usize,
|
|
nonlinearity: NonLinearity,
|
|
) -> Result<()> {
|
|
let gain = nonlinearity.gain();
|
|
let std = gain / (fan_in as f32).sqrt();
|
|
normal(tensor, 0.0, std)
|
|
}
|
|
|
|
/// Xavier uniform initialization
|
|
pub fn xavier_uniform(tensor: &mut Tensor, fan_in: usize, fan_out: usize) -> Result<()> {
|
|
let bound = (6.0 / (fan_in + fan_out) as f32).sqrt();
|
|
uniform(tensor, -bound, bound)
|
|
}
|
|
|
|
/// Xavier normal initialization
|
|
pub fn xavier_normal(tensor: &mut Tensor, fan_in: usize, fan_out: usize) -> Result<()> {
|
|
let std = (2.0 / (fan_in + fan_out) as f32).sqrt();
|
|
normal(tensor, 0.0, std)
|
|
}
|