Source browser

examples/shaders/cnn_infer.cpp

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

1// 8x8 grayscale -> two valid 3x3 filters -> ReLU -> global mean -> sigmoid classifier.
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 Parameters { // Read the complete flat learned parameter vector through buffer binding 1.
10    float parameters[]; // Store the complete flat learned parameter vector.
11}; // Finish this declaration.
12layout(std430, binding = 2) writeonly buffer Predictions { // Write one predicted class probability per sample through buffer binding 2.
13    float predictions[]; // Store one predicted class probability per sample.
14}; // Finish this declaration.
15uniform uint batch_count; // Set the number of independent samples in this model batch.
16float logistic(float value) { // Evaluate the sigmoid with a sign-selected, numerically stable exponential.
17    if (value >= 0.0) // Use the negative exponent when the input logit is nonnegative.
18        return 1.0 / (1.0 + exp(-value)); // Convert a nonnegative logit into its sigmoid probability.
19    float e = exp(value); // Use a nonpositive exponent for the negative logit branch.
20    return e / (1.0 + e); // Convert a negative logit into its sigmoid probability.
21} // End this block.
22void main() { // Evaluate the two-filter convolutional classifier for one image.
23    uint row = gl_GlobalInvocationID.x; // Select the independent input sample assigned to this invocation.
24    if (row >= batch_count) // Ignore invocations outside the sample batch.
25        return; // Stop this invocation before performing another resource access.
26    float logit = parameters[22]; // Initialize the binary output logit from its learned bias.
27    for (uint channel = 0u; channel < 2u; ++channel) { // Evaluate each of the two learned convolution filters.
28        float pooled = 0.0; // Start the mean feature and its derivative at zero.
29        for (uint y = 0u; y < 6u; ++y) // Visit the six valid output rows of the three-by-three filter.
30            for (uint x = 0u; x < 6u; ++x) { // Visit the six valid output columns of this filter.
31                float activation = parameters[18u + channel]; // Start this convolution window with the filter's learned bias.
32                for (uint ky = 0u; ky < 3u; ++ky) // Visit each row of the three-by-three learned kernel.
33                    for (uint kx = 0u; kx < 3u; ++kx) // Visit each column of the three-by-three learned kernel.
34                        activation += parameters[9u * channel + 3u * ky + kx] * // Multiply this kernel coefficient by its matching input image sample.
35                                      inputs[64u * row + 8u * (y + ky) + x + kx]; // Read the sample from this image's flattened eight-by-eight plane.
36                pooled += max(0.0, activation) / 36.0; // Apply ReLU and add one thirty-sixth of this valid convolution feature.
37            } // End this block.
38        logit += parameters[20u + channel] * pooled; // Connect this mean convolution feature through its learned head weight.
39    } // End this block.
40    predictions[row] = logistic(logit); // Convert the two-filter output logit into its binary class probability.
41} // End this block.