/
khlevnov
/
twinkle
Обзор
Документация
Войти
/
khlevnov
/
twinkle
Код
Запросы
0
Пакеты
0
Релизы
0
CI/CD
Аналитика
Безопасность
master
src/nn/linear.rs
55 строк
1 KB
Oleg Khlevnov
wip
09 апр 2024, 19:49
09 апр 2024, 19:49
bd44fe2
Код
Авторство
О чём код?
use num_traits::{Num, NumCast}; use crate::nn::init::uniform; use crate::nn::module::Module; use crate::Tensor; pub struct Linear<T> { weight: Tensor<T>, bias: Option<Tensor<T>>, } impl<T> Linear<T> where T: Copy + Num + NumCast + 'static, { pub fn new(in_features: usize, out_features: usize, bias: bool) -> Linear<T> { let k = 1. / in_features as f64; let weight = uniform(&[in_features, out_features], -k.sqrt(), k.sqrt()); weight.requires_grad_(true); let bias = if bias { let bias = uniform(&[out_features], -k.sqrt(), k.sqrt()); bias.requires_grad_(true); Some(bias) } else { None }; Self { weight, bias, } } } impl<T> Module<T> for Linear<T> where T: Copy + Num + 'static, { fn forward(&self, tensor: Tensor<T>) -> Tensor<T> { match &self.bias { Some(bias) => &tensor.mm(&self.weight) + bias, None => tensor.mm(&self.weight), } } fn parameters(&self) -> Vec<Tensor<T>> { let mut params = vec![self.weight.clone()]; if let Some(bias) = &self.bias { params.push(bias.clone()) } params } }