Source browser

examples/shaders/cnn_gradient.cpp

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

1// Train convolution coefficients, biases, and the classifier head with analytic gradients.
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 image loss with respect to one CNN parameter.
26    uint index = gl_GlobalInvocationID.x, row = index / 23u, parameter = index % 23u; // Select one sample-major derivative among all 23 CNN parameters.
27    if (row >= batch_count) // Ignore invocations outside the sample batch.
28        return; // Stop this invocation before performing another resource access.
29    float logit = parameters[22], derivative = parameter == 22u ? 1.0 : 0.0; // Initialize the output bias and its derivative for this parameter.
30    for (uint channel = 0u; channel < 2u; ++channel) { // Evaluate each of the two learned convolution filters.
31        float pooled = 0.0, pooled_derivative = 0.0; // Start the mean feature and its derivative at zero.
32        for (uint y = 0u; y < 6u; ++y) // Visit the six valid output rows of the three-by-three filter.
33            for (uint x = 0u; x < 6u; ++x) { // Visit the six valid output columns of this filter.
34                float activation = parameters[18u + channel]; // Start this convolution window with the filter's learned bias.
35                for (uint ky = 0u; ky < 3u; ++ky) // Visit each row of the three-by-three learned kernel.
36                    for (uint kx = 0u; kx < 3u; ++kx) // Visit each column of the three-by-three learned kernel.
37                        activation += parameters[9u * channel + 3u * ky + kx] * // Multiply this kernel coefficient by its matching input image sample.
38                                      inputs[64u * row + 8u * (y + ky) + x + kx]; // Read the sample from this image's flattened eight-by-eight plane.
39                if (activation > 0.0) { // Pass feature derivatives only through the active ReLU branch.
40                    pooled += activation / 36.0; // Accumulate the active convolution feature in the global mean.
41                    if (parameter < 18u && parameter / 9u == channel) { // Differentiate a coefficient only when it belongs to this filter.
42                        uint coefficient = parameter % 9u; // Locate the selected three-by-three coefficient within its filter.
43                        pooled_derivative += // Accumulate this kernel coefficient's contribution to the pooled derivative.
44                            inputs[64u * row + 8u * (y + coefficient / 3u) + x + coefficient % 3u] / // Read the sample from this image's flattened eight-by-eight plane.
45                            36.0; // Normalize the derivative by the valid six-by-six output area.
46                    } else if (parameter == 18u + channel) // Handle the derivative of this filter's additive bias.
47                        pooled_derivative += 1.0 / 36.0; // Normalize the derivative by the valid six-by-six output area.
48                } // End this block.
49            } // End this block.
50        logit += parameters[20u + channel] * pooled; // Connect this mean convolution feature through its learned head weight.
51        derivative += parameters[20u + channel] * pooled_derivative; // Propagate the pooled feature derivative through the head weight.
52        if (parameter == 20u + channel) // Handle the derivative of the head weight for this filter.
53            derivative += pooled; // The head weight derivative is this filter's pooled feature.
54    } // End this block.
55    gradients[index] = (logistic(logit) - targets[row]) * derivative; // Multiply the output-loss derivative by the full parameter chain-rule derivative.
56} // End this block.