Source browser

examples/shaders/ann_gradient.cpp

Read the complete file below or download its original bytes. Line numbers provide stable links to a declaration or processing step.

1// Analytic binary cross-entropy gradients for every row and every ANN parameter.
2#ifdef __cplusplus // Expose host declarations only when a C++ compiler reads this shader.
3#include "shdr.h" // Include shdr.h declarations.
4#endif // Finish this conditional compilation branch.
5layout(local_size_x = 64) in; // Group 64 independent invocations along x.
6layout(std430, binding = 0) readonly buffer Inputs { // Read the model's sample-major input features through buffer binding 0.
7    float inputs[]; // Store the model's sample-major input features.
8}; // Finish this declaration.
9layout(std430, binding = 1) readonly buffer Targets { // Read one binary training label per sample through buffer binding 1.
10    float targets[]; // Store one binary training label per sample.
11}; // Finish this declaration.
12layout(std430, binding = 2) readonly buffer Parameters { // Read the complete flat learned parameter vector through buffer binding 2.
13    float parameters[]; // Store the complete flat learned parameter vector.
14}; // Finish this declaration.
15layout(std430, binding = 3) writeonly buffer Gradients { // Write sample-major derivatives for every trainable parameter through buffer binding 3.
16    float gradients[]; // Store sample-major derivatives for every trainable parameter.
17}; // Finish this declaration.
18uniform uint batch_count; // Set the number of independent samples in this model batch.
19float logistic(float value) { // Evaluate the sigmoid with a sign-selected, numerically stable exponential.
20    if (value >= 0.0) // Use the negative exponent when the input logit is nonnegative.
21        return 1.0 / (1.0 + exp(-value)); // Convert a nonnegative logit into its sigmoid probability.
22    float e = exp(value); // Use a nonpositive exponent for the negative logit branch.
23    return e / (1.0 + e); // Convert a negative logit into its sigmoid probability.
24} // End this block.
25void main() { // Differentiate one sample loss with respect to one ANN parameter.
26    uint index = gl_GlobalInvocationID.x, row = index / 33u, parameter = index % 33u; // Select one sample-major derivative among all 33 ANN parameters.
27    if (row >= batch_count) // Ignore invocations outside the sample batch.
28        return; // Stop this invocation before performing another resource access.
29    float x = inputs[2u * row], y = inputs[2u * row + 1u], logit = parameters[32]; // Read the two input features and initialize the output with its bias.
30    float activations[8], preactivations[8]; // Retain hidden preactivations and ReLU outputs for the chain rule.
31    for (uint neuron = 0u; neuron < 8u; ++neuron) { // Evaluate each of the eight hidden neurons.
32        preactivations[neuron] = parameters[2u * neuron] * x + parameters[2u * neuron + 1u] * y + // Multiply the two input features by this neuron's learned weights.
33                                 parameters[16u + neuron]; // Include this hidden neuron's bias before the ReLU.
34        activations[neuron] = max(0.0, preactivations[neuron]); // Form the hidden neuron's weighted input before its nonlinearity.
35        logit += activations[neuron] * parameters[24u + neuron]; // Keep positive preactivations and set the negative branch to zero.
36    } // End this block.
37    float delta = logistic(logit) - targets[row], gradient = delta; // Use prediction minus target as the cross-entropy derivative with respect to the logit.
38    if (parameter < 16u) { // Differentiate one of the sixteen hidden input weights.
39        uint neuron = parameter / 2u; // Find the hidden neuron owning this input weight.
40        gradient = preactivations[neuron] > 0.0 // Form the hidden neuron's weighted input before its nonlinearity.
41                       ? delta * parameters[24u + neuron] * (parameter % 2u == 0u ? x : y) // Multiply the output derivative by its output weight and the selected input feature.
42                       : 0.0; // Use zero derivative for the inactive ReLU branch.
43    } else if (parameter < 24u) { // Differentiate one of the eight hidden biases.
44        uint neuron = parameter - 16u; // Find the hidden neuron owning this bias.
45        gradient = preactivations[neuron] > 0.0 ? delta * parameters[24u + neuron] : 0.0; // Form the hidden neuron's weighted input before its nonlinearity.
46    } else if (parameter < 32u) // Differentiate one of the eight output weights.
47        gradient = delta * activations[parameter - 24u]; // Multiply the output logit derivative by this hidden activation.
48    gradients[index] = gradient; // Store this sample's derivative for the following batch SGD stage.
49} // End this block.