brain2qwerty 0.0.1

Brain2Qwerty V1/V2 MEG neural decoding inference in Rust (parity-tested vs Python)
Documentation 
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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! SimpleConv encoder (neuraltrain SimpleConvModel + ConvSequence).

use crate::config::SimpleConvConfig;
use crate::model::channel_merger::ChannelMerger;
use crate::model::subject_layers::SubjectLayers;
use crate::tensor::Tensor;
use crate::weights::ParamBuf;

pub struct ConvBlock {
    pub conv_w: Tensor,
    pub conv_b: Tensor,
    pub padding: usize,
    pub dilation: usize,
    pub bn_w: Option<Tensor>,
    pub bn_b: Option<Tensor>,
    pub bn_running_mean: Option<Tensor>,
    pub bn_running_var: Option<Tensor>,
    pub layer_scale: Option<Tensor>,
    pub use_gelu: bool,
    pub leakiness: f32,
    pub use_activation: bool,
    pub skip: bool,
}

pub struct SimpleConvEncoder {
    pub merger: Option<ChannelMerger>,
    pub initial_linear: Option<(Tensor, Tensor)>,
    pub subject_layers: Option<SubjectLayers>,
    pub blocks: Vec<ConvBlock>,
    pub out_channels: usize,
}

impl SimpleConvEncoder {
    pub fn forward(&self, x: &Tensor, subject_ids: &[usize], positions: Option<&Tensor>) -> Tensor {
        let length = x.shape[2];
        let mut z = x.clone();

        if let Some(ref merger) = self.merger {
            let pos = positions.expect("channel_positions required for merger");
            z = merger.forward(&z, subject_ids, pos);
        }

        if let Some((ref w, ref b)) = self.initial_linear {
            z = z.conv1d(w, Some(b), 1, 0, 1, 1);
        }

        if let Some(ref sl) = self.subject_layers {
            z = sl.forward(&z, subject_ids);
        }

        for block in &self.blocks {
            let old = z.clone();
            z = z.conv1d(
                &block.conv_w,
                Some(&block.conv_b),
                1,
                block.padding,
                block.dilation,
                1,
            );
            if let (Some(bn_w), Some(bn_b), Some(rm), Some(rv)) = (
                &block.bn_w,
                &block.bn_b,
                &block.bn_running_mean,
                &block.bn_running_var,
            ) {
                z = z.batch_norm1d_eval(bn_w, bn_b, rm, rv, 1e-5);
            }
            if block.use_activation {
                z = if block.use_gelu {
                    z.gelu()
                } else if block.leakiness > 0.0 {
                    z.leaky_relu(block.leakiness)
                } else {
                    z.relu()
                };
            }
            if let Some(ref scale) = block.layer_scale {
                let boost = 5.0f32;
                let (b, c, t) = (z.shape[0], z.shape[1], z.shape[2]);
                for bi in 0..b {
                    for ci in 0..c {
                        let s = scale.data[ci] * boost;
                        for ti in 0..t {
                            let idx = (bi * c + ci) * t + ti;
                            z.data[idx] *= s;
                        }
                    }
                }
            }
            if block.skip && old.shape == z.shape {
                z = z.add(&old);
            }
        }

        assert!(z.shape[2] >= length);
        z.narrow(2, 0, length)
    }
}

pub fn build_channel_sizes(
    cfg: &SimpleConvConfig,
    in_channels: usize,
    out_channels: usize,
) -> Vec<usize> {
    let mut sizes = vec![in_channels];
    for _k in 0..cfg.depth {
        sizes.push(cfg.hidden);
    }
    sizes.push(out_channels);
    if sizes.len() > cfg.depth + 1 {
        let last = sizes.len() - 1;
        sizes[last] = out_channels;
    }
    sizes
}

pub fn layer_scale_from_param(p: &ParamBuf, boost: f32) -> Tensor {
    let mut data = p.data.clone();
    for v in &mut data {
        *v = (*v / boost) * boost;
    }
    Tensor::from_vec(data, p.shape.clone())
}