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 →
2.6 KiB73 linesraw
1module mlp;
2
3import common;
4
5__include mlvec;
6
7// We use Float16 for the CoopVec component type since it is more widely supported.
8//
9static const CoopVecComponentType kComponentType = CoopVecComponentType.Float16;
10
11public struct FeedForwardLayer<int InputSize, int OutputSize>
12{
13    internal void* weights;
14    internal void* weightsGrad;
15    internal void* biases;
16    internal void* biasesGrad;
17
18    public MLVec<OutputSize> eval(MLVec<InputSize> input)
19    {
20        // Compute mul(weights, inputVec) + biases.
21        // `weights` is treated as an OutputSize(row) x InputSize(col) matrix.
22        var output = coopVecMatMulAdd<NFloat, OutputSize>(
23            input.data, kComponentType, // input and format
24            weights, kComponentType, // weights and format
25            biases, kComponentType, // biases and format
26            CoopVecMatrixLayout.RowMajor, // matrix layout
27            false, // transpose matrix? must be `false` since we specified RowMajor.
28            InputSize * sizeof(NFloat)); // matrix stride
29        output = max(output, output * 0.001h); // Leaky ReLU activation
30        return {output};
31    }
32
33    [BackwardDerivativeOf(eval)]
34    public void evalBwd(
35        inout DifferentialPair<MLVec<InputSize>> input,
36        MLVec<OutputSize> resultGrad)
37    {
38        let fwd = eval(input.p);
39
40        // Back-prop resultGrad through activation.
41        [ForceUnroll]
42        for (int i = 0; i < OutputSize; i++)
43        {
44            if (fwd.data[i] < 0.0)
45                resultGrad.data[i] *= 0.01h;
46        }
47
48        // Back-prop gradients to the weights matrix.
49        coopVecOuterProductAccumulate(
50            resultGrad.data,
51            input.p.data,
52            weightsGrad,
53            0, // matrixStride, ignored since layout is TrainingOptimal
54            CoopVecMatrixLayout.TrainingOptimal, // matrix layout, must be TrainingOptimal.
55            kComponentType);
56        
57        // Back-prop gradients to the biases vector.
58        coopVecReduceSumAccumulate(resultGrad.data, (void*)biasesGrad);
59
60        // Back-prop gradients to the input vector by computing
61        // mul(transpose(weights), resultGrad).
62        // By specifying the matrix layout as ColumnMajor, we can
63        // achieve the effect of transposing the weights matrix.
64        let dInput = coopVecMatMul<NFloat, InputSize>(
65            resultGrad.data, kComponentType,
66            weights, kComponentType,
67            CoopVecMatrixLayout.ColumnMajor,
68            false,  // transpose, must be `false` since we specified ColumnMajor.
69            InputSize * sizeof(NFloat));
70
71        input = {input.p, {dInput}};
72    }
73}