tch 0.0.1

PyTorch wrappers for rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

use crate::tensor::Tensor;

pub struct BatchNorm2DConfig {
    cudnn_enabled: bool,
    eps: f64,
    momentum: f64,
}

impl Default for BatchNorm2DConfig {
    fn default() -> Self {
        BatchNorm2DConfig {
            cudnn_enabled: true,
            eps: 1e-5,
            momentum: 0.1,
        }
    }
}

pub struct BatchNorm2D {
    config: BatchNorm2DConfig,
    running_mean: Tensor,
    running_var: Tensor,
    ws: Tensor,
    bs: Tensor,
}

impl BatchNorm2D {
    pub fn new(
        vs: &super::var_store::Path,
        out_dim: i64,
        config: BatchNorm2DConfig,
    ) -> BatchNorm2D {
        BatchNorm2D {
            config,
            running_mean: vs.zeros_no_train("running_mean", &[out_dim]),
            running_var: vs.ones_no_train("running_var", &[out_dim]),
            ws: vs.uniform("weight", &[out_dim], 0.0, 1.0),
            bs: vs.zeros("bias", &[out_dim]),
        }
    }
}

impl super::module::ModuleT for BatchNorm2D {
    fn forward_t(&self, xs: &Tensor, train: bool) -> Tensor {
        Tensor::batch_norm(
            xs,
            Some(&self.ws),
            Some(&self.bs),
            Some(&self.running_mean),
            Some(&self.running_var),
            train,
            self.config.momentum,
            self.config.eps,
            self.config.cudnn_enabled,
        )
    }
}