yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Yong HeAdd MLP training examples. (#7550)f28f67d98

master
branchesmasterall refs →
1.5 KiB59 linesraw
1module mlp_sw;
2
3import common;
4
5__include mlvec_sw;
6
7public struct FeedForwardLayer<int InputSize, int OutputSize>
8{
9    public NFloat* weights;
10    public NFloat* weightsGrad;
11    public NFloat* biases;
12    public NFloat* biasesGrad;
13
14    [BackwardDerivative(evalBwd)]
15    public MLVec<OutputSize> eval(MLVec<InputSize> input)
16    {
17        var output = matMulAdd<OutputSize>(
18            input,
19            weights,
20            biases);
21        // ReLU activation
22        for (int i = 0; i < OutputSize; i++)
23            if (output.data[i] < 0.0)
24                output.data[i] *= 0.001h;
25        return output; 
26    }
27
28    public void evalBwd(
29        inout DifferentialPair<MLVec<InputSize>> input,
30        MLVec<OutputSize> resultGrad)
31    {
32        let fwd = eval(input.p);
33
34        // Back-prop resultGrad through activation.
35        for (int i = 0; i < OutputSize; i++)
36        {
37            if (fwd.data[i] < 0.0)
38                resultGrad.data[i] *= 0.01h;
39        }
40
41        // Back-prop gradients to the weights matrix.
42        outerProductAccumulate(
43            resultGrad,
44            input.p,
45            weightsGrad);
46
47        // Back-prop gradients to the biases vector.
48        for (int i = 0; i < OutputSize; i++)
49        {
50            NFloat originalValue;
51            InterlockedAddF16Emulated(biasesGrad + i, resultGrad.data[i], originalValue);
52        }
53
54        // Back-prop gradients to the input vector.
55        let dInput = matMulTransposed<InputSize>(resultGrad, weights);
56
57        input = {input.p, dInput};
58    }
59}