1 files changed, 104 insertions, 78 deletions

diff --git a/src/lib/src/train.c b/src/lib/src/train.c
index dc93f0f..98f58ad 100644
--- a/src/lib/src/train.c
+++ b/src/lib/src/train.c
@@ -38,7 +38,7 @@ typedef struct nnSigmoidGradientElements {
 /// each layer. A data type is defined for these because we allocate all the
 /// required memory up front before entering the training loop.
 typedef struct nnGradientElements {
-  nnActivation type;
+  nnLayerType type;
   // Gradient vector, same size as the layer.
   // This will contain the gradient expression except for the output value of
   // the previous layer.
@@ -57,10 +57,27 @@ void nnInitNet(
   mt19937_64_init(&rng, seed);
 
   for (int l = 0; l < net->num_layers; ++l) {
-    nnMatrix* weights = &net->weights[l];
-    nnMatrix* biases  = &net->biases[l];
+    // Get the layer's weights and biases, if any.
+    nnMatrix* weights = 0;
+    nnMatrix* biases  = 0;
+    switch (net->layers[l].type) {
+    case nnLinear: {
+      nnLinearImpl* linear = &net->layers[l].linear;
+
+      weights = &linear->weights;
+      biases  = &linear->biases;
+      break;
+    }
+    // Activations.
+    case nnRelu:
+    case nnSigmoid:
+      break;
+    }
+    if (!weights || !biases) {
+      continue;
+    }
 
-    const R layer_size = (R)nnLayerInputSize(weights);
+    const R layer_size = (R)nnLayerInputSize(net, l);
     const R scale      = 1. / layer_size;
     const R stdev      = 1. / sqrt((R)layer_size);
     const R sigma      = stdev * stdev;
@@ -128,9 +145,6 @@ void nnTrain(
   // with one sample at a time.
   nnMatrix* errors = calloc(net->num_layers, sizeof(nnMatrix));
 
-  // Allocate the weight transpose matrices up front for backpropagation.
-  // nnMatrix* weights_T = calloc(net->num_layers, sizeof(nnMatrix));
-
   // Allocate the weight delta matrices.
   nnMatrix* weight_deltas = calloc(net->num_layers, sizeof(nnMatrix));
 
@@ -144,30 +158,24 @@ void nnTrain(
   nnMatrix* outputs_T = calloc(net->num_layers, sizeof(nnMatrix));
 
   assert(errors != 0);
-  // assert(weights_T != 0);
   assert(weight_deltas != 0);
   assert(gradient_elems);
   assert(outputs_T);
 
   for (int l = 0; l < net->num_layers; ++l) {
-    const nnMatrix*    layer_weights     = &net->weights[l];
-    const int          layer_output_size = net->weights[l].cols;
-    const nnActivation activation        = net->activations[l];
-
-    errors[l] = nnMatrixMake(1, layer_weights->cols);
-
-    // weights_T[l] = nnMatrixMake(layer_weights->cols, layer_weights->rows);
-    // nnMatrixTranspose(layer_weights, &weights_T[l]);
-
-    weight_deltas[l] = nnMatrixMake(layer_weights->rows, layer_weights->cols);
+    const int          layer_input_size  = nnLayerInputSize(net, l);
+    const int          layer_output_size = nnLayerOutputSize(net, l);
+    const nnLayerImpl* layer             = &net->layers[l];
 
-    outputs_T[l] = nnMatrixMake(layer_output_size, 1);
+    errors[l]        = nnMatrixMake(1, layer_output_size);
+    weight_deltas[l] = nnMatrixMake(layer_input_size, layer_output_size);
+    outputs_T[l]     = nnMatrixMake(layer_output_size, 1);
 
     // Allocate the gradient elements and vectors for weight delta calculation.
     nnGradientElements* elems = &gradient_elems[l];
-    elems->type               = activation;
-    switch (activation) {
-    case nnIdentity:
+    elems->type               = layer->type;
+    switch (layer->type) {
+    case nnLinear:
       break; // Gradient vector will be borrowed, no need to allocate.
 
     case nnSigmoid:
@@ -208,6 +216,7 @@ void nnTrain(
 
     // For now, we train with one sample at a time.
     for (int sample = 0; sample < inputs->rows; ++sample) {
+      // TODO: Introduce a BorrowMut.
       // Slice the input and target matrices with the batch size.
       // We are not mutating the inputs, but we need the cast to borrow.
       nnMatrix training_inputs =
@@ -219,15 +228,16 @@ void nnTrain(
       // Assuming one training input per iteration for now.
       nnMatrixTranspose(&training_inputs, &training_inputs_T);
 
-      // Run a forward pass and compute the output layer error relevant to the
-      // derivative: o-t.
-      //   Error: (t-o)^2
-      //   dE/do = -2(t-o)
-      //         = +2(o-t)
+      // Forward pass.
+      nnQuery(net, query, &training_inputs);
+
+      // Compute the error derivative: o-t.
+      //   Error: 1/2 (t-o)^2
+      //   dE/do = -(t-o)
+      //         = +(o-t)
       // Note that we compute o-t instead to remove that outer negative sign.
       // The 2 is dropped because we are only interested in the direction of the
       // gradient. The learning rate controls the magnitude.
-      nnQuery(net, query, &training_inputs);
       nnMatrixSub(
           training_outputs, &training_targets, &errors[net->num_layers - 1]);
 
@@ -236,68 +246,86 @@ void nnTrain(
         nnMatrixTranspose(&query->layer_outputs[l], &outputs_T[l]);
       }
 
-      // Update weights and biases for each internal layer, backpropagating
+      // Update weights and biases for each internal layer, back-propagating
       // errors along the way.
       for (int l = net->num_layers - 1; l >= 0; --l) {
-        const nnMatrix*     layer_output  = &query->layer_outputs[l];
-        nnMatrix*           layer_weights = &net->weights[l];
-        nnMatrix*           layer_biases  = &net->biases[l];
-        nnGradientElements* elems         = &gradient_elems[l];
-        nnMatrix*           gradient      = &elems->gradient;
-        const nnActivation  activation    = net->activations[l];
-
-        // Compute the gradient (the part of the expression that does not
-        // contain the output of the previous layer).
+        const nnMatrix*     layer_output = &query->layer_outputs[l];
+        nnGradientElements* elems        = &gradient_elems[l];
+        nnMatrix*           gradient     = &elems->gradient;
+        nnLayerImpl*        layer        = &net->layers[l];
+
+        // Compute this layer's gradient.
+        //
+        // By "gradient" we mean the expression common to the weights and bias
+        // gradients. This is the part of the expression that does not contain
+        // this layer's input.
         //
-        // Identity: G = error_k
-        // Sigmoid:  G = error_k * output_k * (1 - output_k).
-        // Relu:     G = error_k * (output_k > 0 ? 1 : 0)
-        switch (activation) {
-        case nnIdentity:
+        // Linear:   G = id
+        // Relu:     G = (output_k > 0 ? 1 : 0)
+        // Sigmoid:  G = output_k * (1 - output_k)
+        switch (layer->type) {
+        case nnLinear: {
           // TODO: Just copy the pointer?
           *gradient = nnMatrixBorrow(&errors[l]);
           break;
+        }
+        case nnRelu:
+          nnMatrixGt(layer_output, 0, gradient);
+          break;
         case nnSigmoid:
           nnMatrixSub(&elems->sigmoid.ones, layer_output, gradient);
           nnMatrixMulPairs(layer_output, gradient, gradient);
-          nnMatrixMulPairs(&errors[l], gradient, gradient);
-          break;
-        case nnRelu:
-          nnMatrixGt(layer_output, 0, gradient);
-          nnMatrixMulPairs(&errors[l], gradient, gradient);
           break;
         }
 
-        // Outer product to compute the weight deltas.
-        const nnMatrix* output_T =
-            (l == 0) ? &training_inputs_T : &outputs_T[l - 1];
-        nnMatrixMul(output_T, gradient, &weight_deltas[l]);
-
-        // Backpropagate the error before updating weights.
+        // Back-propagate the error.
+        //
+        // This combines this layer's gradient with the back-propagated error,
+        // which is the combination of the gradients of subsequent layers down
+        // to the output layer error.
+        //
+        // Note that this step uses the layer's original weights.
         if (l > 0) {
-          // G * W^T == G *^T W.
-          // nnMatrixMul(gradient, &weights_T[l], &errors[l-1]);
-          nnMatrixMulRows(gradient, layer_weights, &errors[l - 1]);
+          switch (layer->type) {
+          case nnLinear: {
+            const nnMatrix* layer_weights = &layer->linear.weights;
+            // E * W^T == E *^T W.
+            // Using nnMatrixMulRows, we avoid having to transpose the weight
+            // matrix.
+            nnMatrixMulRows(&errors[l], layer_weights, &errors[l - 1]);
+            break;
+          }
+          // For activations, the error back-propagates as is but multiplied by
+          // the layer's gradient.
+          case nnRelu:
+          case nnSigmoid:
+            nnMatrixMulPairs(&errors[l], gradient, &errors[l - 1]);
+            break;
+          }
         }
 
-        // Update weights.
-        nnMatrixScale(&weight_deltas[l], params->learning_rate);
-        // The gradient has a negative sign from -(t - o), but we have computed
-        // e = o - t instead, so we can subtract directly.
-        // nnMatrixAdd(layer_weights, &weight_deltas[l], layer_weights);
-        nnMatrixSub(layer_weights, &weight_deltas[l], layer_weights);
-
-        // Update weight transpose matrix for the next training iteration.
-        // nnMatrixTranspose(layer_weights, &weights_T[l]);
-
-        // Update biases.
-        // This is the same formula as for weights, except that the o_j term is
-        // just 1. We can simply re-use the gradient that we have already
-        // computed for the weight update.
-        // nnMatrixMulAdd(layer_biases, gradient, params->learning_rate,
-        // layer_biases);
-        nnMatrixMulSub(
-            layer_biases, gradient, params->learning_rate, layer_biases);
+        // Update layer weights.
+        if (layer->type == nnLinear) {
+          nnLinearImpl* linear        = &layer->linear;
+          nnMatrix*     layer_weights = &linear->weights;
+          nnMatrix*     layer_biases  = &linear->biases;
+
+          // Outer product to compute the weight deltas.
+          // This layer's input is the previous layer's output.
+          const nnMatrix* input_T =
+              (l == 0) ? &training_inputs_T : &outputs_T[l - 1];
+          nnMatrixMul(input_T, gradient, &weight_deltas[l]);
+
+          // Update weights.
+          nnMatrixScale(&weight_deltas[l], params->learning_rate);
+          nnMatrixSub(layer_weights, &weight_deltas[l], layer_weights);
+
+          // Update biases.
+          // This is the same formula as for weights, except that the o_j term
+          // is just 1.
+          nnMatrixMulSub(
+              layer_biases, gradient, params->learning_rate, layer_biases);
+        }
       }
 
       // TODO: Add this under a verbose debugging mode.
@@ -334,12 +362,11 @@ void nnTrain(
   for (int l = 0; l < net->num_layers; ++l) {
     nnMatrixDel(&errors[l]);
     nnMatrixDel(&outputs_T[l]);
-    // nnMatrixDel(&weights_T[l]);
     nnMatrixDel(&weight_deltas[l]);
 
     nnGradientElements* elems = &gradient_elems[l];
     switch (elems->type) {
-    case nnIdentity:
+    case nnLinear:
       break; // Gradient vector is borrowed, no need to deallocate.
 
     case nnSigmoid:
@@ -355,7 +382,6 @@ void nnTrain(
   nnMatrixDel(&training_inputs_T);
   free(errors);
   free(outputs_T);
-  // free(weights_T);
   free(weight_deltas);
   free(gradient_elems);
 }