Source browser

examples/shaders/nn_sgd.cpp

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

1// Batch-reduce independent parameter gradients and perform an in-place SGD update.
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 Gradients { // Read sample-major derivatives for every trainable parameter through buffer binding 0.
7    float gradients[]; // Store sample-major derivatives for every trainable parameter.
8}; // Finish this declaration.
9layout(std430, binding = 1) buffer Parameters { // Read and update the complete flat learned parameter vector through buffer binding 1.
10    float parameters[]; // Store the complete flat learned parameter vector.
11}; // Finish this declaration.
12uniform uint batch_count; // Set the number of independent samples in this model batch.
13uniform uint parameter_count; // Set the number of trainable values in each sample gradient record.
14uniform float learning_rate; // Scale the averaged loss derivative before updating each parameter.
15uniform float weight_decay; // Add an L2 parameter penalty to the averaged loss derivative.
16void main() { // Average the batch gradients and update one trainable parameter.
17    uint parameter = gl_GlobalInvocationID.x; // Assign one distinct trainable parameter to this invocation.
18    if (parameter >= parameter_count) // Ignore padding outside the complete model vector.
19        return; // Stop this invocation before performing another resource access.
20    float gradient = 0.0; // Start the batch gradient reduction at zero.
21    for (uint row = 0u; row < batch_count; ++row) // Read this parameter's derivative from every sample.
22        gradient += gradients[row * parameter_count + parameter]; // Accumulate the sample-major derivative at this parameter index.
23    parameters[parameter] -= // Update this parameter in place after the gradient stage has completed.
24        learning_rate * (gradient / float(batch_count) + weight_decay * parameters[parameter]); // Update this parameter in place after the gradient stage has completed.
25} // End this block.