FirstVersion

Xenith/shader.comp

@@ -2,35 +2,64 @@
 
 layout(local_size_x = 256) in;
 
-layout(std430, binding = 0) buffer WeightBuffer { float weights[]; };
-layout(std430, binding = 1) buffer BiasBuffer   { float biases[]; };
-layout(std430, binding = 2) buffer OutputBuffer { float outputs[]; };
-layout(std430, binding = 3) buffer ErrorBuffer  { float errors[]; };
+layout(std430, binding = 0) buffer Weights { float W[]; };
+layout(std430, binding = 1) buffer Biases  { float B[]; };
+layout(std430, binding = 2) buffer Outputs { float O[]; };
+layout(std430, binding = 3) buffer Errors  { float E[]; };
+layout(std430, binding = 4) buffer Targets { float T[]; };
 
 layout(push_constant) uniform Params {
-    uint prevLayerSize;
-    uint nextLayerSize;
-    uint weightOffset;
-    uint biasOffset;
-    uint outOffset;
-    uint errOffset;
+    uint mode;     // 0: FF, 1: OutError, 2: BackProp, 3: Update
+    uint prevSize;
+    uint nextSize;
+    uint wOff;
+    uint bOff;
+    uint oOff;
+    uint nextOOff;
     float lr;
 } p;
 
+float sigmoid(float x) { return 1.0 / (1.0 + exp(-x)); }
+float dSigmoid(float x) { return x * (1.0 - x); }
+
 void main() {
-    uint j = gl_GlobalInvocationID.x; // Индекс нейрона следующего слоя
+    uint idx = gl_GlobalInvocationID.x;
 
-    if (j < p.nextLayerSize) {
-        float errorTerm = p.lr * errors[p.errOffset + j];
-
-        // Обновляем веса, входящие в этот нейрон
-        for (uint k = 0; k < p.prevLayerSize; k++) {
-            uint wIdx = p.weightOffset + (j * p.prevLayerSize + k);
-            uint outIdx = p.outOffset + k;
-            weights[wIdx] += errorTerm * outputs[outIdx];
+    // MODE 0: Прямое распространение (Forward Pass)
+    if (p.mode == 0) {
+        if (idx < p.nextSize) {
+            float sum = B[p.bOff + idx];
+            for (uint i = 0; i < p.prevSize; i++) {
+                sum += O[p.oOff + i] * W[p.wOff + idx * p.prevSize + i];
+            }
+            O[p.nextOOff + idx] = sigmoid(sum);
+        }
+    }
+    // MODE 1: Ошибка выходного слоя
+    else if (p.mode == 1) {
+        if (idx < p.nextSize) {
+            float outVal = O[p.nextOOff + idx];
+            E[p.nextOOff + idx] = (T[idx] - outVal) * dSigmoid(outVal);
+        }
+    }
+    // MODE 2: Обратное распространение ошибки (Hidden layers)
+    else if (p.mode == 2) {
+        if (idx < p.prevSize) {
+            float errSum = 0.0;
+            for (uint i = 0; i < p.nextSize; i++) {
+                errSum += E[p.nextOOff + i] * W[p.wOff + i * p.prevSize + idx];
+            }
+            E[p.oOff + idx] = errSum * dSigmoid(O[p.oOff + idx]);
+        }
+    }
+    // MODE 3: Обновление весов и смещений
+    else if (p.mode == 3) {
+        if (idx < p.nextSize) {
+            float errTerm = E[p.nextOOff + idx] * p.lr;
+            for (uint i = 0; i < p.prevSize; i++) {
+                W[p.wOff + idx * p.prevSize + i] += errTerm * O[p.oOff + i];
+            }
+            B[p.bOff + idx] += errTerm;
         }
-
-        // Обновляем биас этого нейрона
-        biases[p.biasOffset + j] += errorTerm;
     }
 }