1 files changed, 114 insertions, 104 deletions

diff --git a/src/lib/src/neuralnet.c b/src/lib/src/neuralnet.c
index a5fc59b..4322b8c 100644
--- a/src/lib/src/neuralnet.c
+++ b/src/lib/src/neuralnet.c
@@ -7,11 +7,65 @@
 #include <assert.h>
 #include <stdlib.h>
 
+static void MakeLayerImpl(
+    int prev_layer_output_size, const nnLayer* layer, nnLayerImpl* impl) {
+  impl->type = layer->type;
+
+  switch (layer->type) {
+  case nnLinear: {
+    const nnLinearParams* params = &layer->linear;
+    nnLinearImpl*         linear = &impl->linear;
+
+    if ((params->input_size > 0) && (params->output_size > 0)) {
+      const int rows  = params->input_size;
+      const int cols  = params->output_size;
+      linear->weights = nnMatrixMake(rows, cols);
+      linear->biases  = nnMatrixMake(1, cols);
+      linear->owned   = true;
+    } else {
+      linear->weights = params->weights;
+      linear->biases  = params->biases;
+      linear->owned   = false;
+    }
+
+    impl->input_size  = linear->weights.rows;
+    impl->output_size = linear->weights.cols;
+
+    break;
+  }
+
+  // Activation layers.
+  case nnRelu:
+  case nnSigmoid:
+    impl->input_size  = prev_layer_output_size;
+    impl->output_size = prev_layer_output_size;
+    break;
+  }
+}
+
+static void DeleteLayer(nnLayerImpl* layer) {
+  switch (layer->type) {
+  case nnLinear: {
+    nnLinearImpl* linear = &layer->linear;
+    if (linear->owned) {
+      nnMatrixDel(&linear->weights);
+      nnMatrixDel(&linear->biases);
+    }
+    break;
+  }
+
+  // No parameters for these layers.
+  case nnRelu:
+  case nnSigmoid:
+    break;
+  }
+}
+
 nnNeuralNetwork* nnMakeNet(
-    int num_layers, const int* layer_sizes, const nnActivation* activations) {
+    const nnLayer* layers, int num_layers, int input_size) {
+  assert(layers);
   assert(num_layers > 0);
-  assert(layer_sizes);
-  assert(activations);
+  assert(input_size > 0);
 
   nnNeuralNetwork* net = calloc(1, sizeof(nnNeuralNetwork));
   if (net == 0) {
@@ -20,84 +74,38 @@ nnNeuralNetwork* nnMakeNet(
 
   net->num_layers = num_layers;
 
-  net->weights     = calloc(num_layers, sizeof(nnMatrix));
-  net->biases      = calloc(num_layers, sizeof(nnMatrix));
-  net->activations = calloc(num_layers, sizeof(nnActivation));
-  if ((net->weights == 0) || (net->biases == 0) || (net->activations == 0)) {
+  net->layers = calloc(num_layers, sizeof(nnLayerImpl));
+  if (net->layers == 0) {
     nnDeleteNet(&net);
     return 0;
   }
 
+  int prev_layer_output_size = input_size;
   for (int l = 0; l < num_layers; ++l) {
-    // layer_sizes = { input layer size, first hidden layer size, ...}
-    const int layer_input_size  = layer_sizes[l];
-    const int layer_output_size = layer_sizes[l + 1];
-
-    // We store the transpose of the weight matrix as written in textbooks.
-    // Our vectors are row vectors and the matrices row-major.
-    const int rows = layer_input_size;
-    const int cols = layer_output_size;
-
-    net->weights[l]     = nnMatrixMake(rows, cols);
-    net->biases[l]      = nnMatrixMake(1, cols);
-    net->activations[l] = activations[l];
+    MakeLayerImpl(prev_layer_output_size, &layers[l], &net->layers[l]);
+    prev_layer_output_size = net->layers[l].output_size;
   }
 
   return net;
 }
 
-void nnDeleteNet(nnNeuralNetwork** net) {
-  if ((!net) || (!(*net))) {
+void nnDeleteNet(nnNeuralNetwork** ppNet) {
+  if ((!ppNet) || (!(*ppNet))) {
     return;
   }
-  if ((*net)->weights != 0) {
-    for (int l = 0; l < (*net)->num_layers; ++l) {
-      nnMatrixDel(&(*net)->weights[l]);
-    }
-    free((*net)->weights);
-    (*net)->weights = 0;
-  }
-  if ((*net)->biases != 0) {
-    for (int l = 0; l < (*net)->num_layers; ++l) {
-      nnMatrixDel(&(*net)->biases[l]);
-    }
-    free((*net)->biases);
-    (*net)->biases = 0;
-  }
-  if ((*net)->activations) {
-    free((*net)->activations);
-    (*net)->activations = 0;
-  }
-  free(*net);
-  *net = 0;
-}
-
-void nnSetWeights(nnNeuralNetwork* net, const R* weights) {
-  assert(net);
-  assert(weights);
+  nnNeuralNetwork* net = *ppNet;
 
   for (int l = 0; l < net->num_layers; ++l) {
-    nnMatrix* layer_weights = &net->weights[l];
-    R*        layer_values  = layer_weights->values;
-
-    for (int j = 0; j < layer_weights->rows * layer_weights->cols; ++j) {
-      *layer_values++ = *weights++;
-    }
+    DeleteLayer(&net->layers[l]);
   }
-}
-
-void nnSetBiases(nnNeuralNetwork* net, const R* biases) {
-  assert(net);
-  assert(biases);
-
-  for (int l = 0; l < net->num_layers; ++l) {
-    nnMatrix* layer_biases = &net->biases[l];
-    R*        layer_values = layer_biases->values;
 
-    for (int j = 0; j < layer_biases->rows * layer_biases->cols; ++j) {
-      *layer_values++ = *biases++;
-    }
+  if (net->layers) {
+    free(net->layers);
+    net->layers = 0;
   }
+
+  free(net);
+  *ppNet = 0;
 }
 
 void nnQuery(
@@ -114,35 +122,40 @@ void nnQuery(
     nnMatrix input_vector = nnMatrixBorrowRows((nnMatrix*)input, i, 1);
 
     for (int l = 0; l < net->num_layers; ++l) {
-      const nnMatrix* layer_weights = &net->weights[l];
-      const nnMatrix* layer_biases  = &net->biases[l];
-      // Y^T = (W*X)^T = X^T*W^T
-      //
-      // TODO: If we had a row-row matrix multiplication, we could compute:
-      //   Y^T = W ** X^T
-      // The row-row multiplication could be more cache-friendly. We just need
-      // to store W as is, without transposing.
-      // We could also rewrite the original Mul function to go row x row,
-      // decomposing the multiplication. Preserving the original meaning of Mul
-      // makes everything clearer.
       nnMatrix output_vector =
           nnMatrixBorrowRows(&query->layer_outputs[l], i, 1);
-      nnMatrixMul(&input_vector, layer_weights, &output_vector);
-      nnMatrixAddRow(&output_vector, layer_biases, &output_vector);
 
-      switch (net->activations[l]) {
-      case nnIdentity:
-        break; // Nothing to do for the identity function.
-      case nnSigmoid:
-        sigmoid_array(
-            output_vector.values, output_vector.values, output_vector.cols);
+      switch (net->layers[l].type) {
+      case nnLinear: {
+        const nnLinearImpl* linear        = &net->layers[l].linear;
+        const nnMatrix*     layer_weights = &linear->weights;
+        const nnMatrix*     layer_biases  = &linear->biases;
+
+        // Y^T = (W*X)^T = X^T*W^T
+        //
+        // TODO: If we had a row-row matrix multiplication, we could compute:
+        //   Y^T = W ** X^T
+        //
+        // The row-row multiplication could be more cache-friendly. We just need
+        // to store W as is, without transposing.
+        //
+        // We could also rewrite the original Mul function to go row x row,
+        // decomposing the multiplication. Preserving the original meaning of
+        // Mul makes everything clearer.
+        nnMatrixMul(&input_vector, layer_weights, &output_vector);
+        nnMatrixAddRow(&output_vector, layer_biases, &output_vector);
         break;
+      }
       case nnRelu:
+        assert(input_vector.cols == output_vector.cols);
         relu_array(
-            output_vector.values, output_vector.values, output_vector.cols);
+            input_vector.values, output_vector.values, output_vector.cols);
+        break;
+      case nnSigmoid:
+        assert(input_vector.cols == output_vector.cols);
+        sigmoid_array(
+            input_vector.values, output_vector.values, output_vector.cols);
         break;
-      default:
-        assert(0);
       }
 
       input_vector = output_vector; // Borrow.
@@ -159,15 +172,15 @@ void nnQueryArray(
   assert(output);
   assert(net->num_layers > 0);
 
-  nnMatrix input_vector = nnMatrixMake(net->weights[0].cols, 1);
+  nnMatrix input_vector = nnMatrixMake(1, nnNetInputSize(net));
   nnMatrixInit(&input_vector, input);
   nnQuery(net, query, &input_vector);
   nnMatrixRowToArray(query->network_outputs, 0, output);
 }
 
-nnQueryObject* nnMakeQueryObject(const nnNeuralNetwork* net, int num_inputs) {
+nnQueryObject* nnMakeQueryObject(const nnNeuralNetwork* net, int batch_size) {
   assert(net);
-  assert(num_inputs > 0);
+  assert(batch_size > 0);
   assert(net->num_layers > 0);
 
   nnQueryObject* query = calloc(1, sizeof(nnQueryObject));
@@ -183,11 +196,12 @@ nnQueryObject* nnMakeQueryObject(const nnNeuralNetwork* net, int num_inputs) {
     free(query);
     return 0;
   }
+
   for (int l = 0; l < net->num_layers; ++l) {
-    const nnMatrix* layer_weights     = &net->weights[l];
-    const int       layer_output_size = nnLayerOutputSize(layer_weights);
-    query->layer_outputs[l] = nnMatrixMake(num_inputs, layer_output_size);
+    const int layer_output_size = nnLayerOutputSize(net, l);
+    query->layer_outputs[l]     = nnMatrixMake(batch_size, layer_output_size);
   }
+
   query->network_outputs = &query->layer_outputs[net->num_layers - 1];
 
   return query;
@@ -213,23 +227,19 @@ const nnMatrix* nnNetOutputs(const nnQueryObject* query) {
 }
 
 int nnNetInputSize(const nnNeuralNetwork* net) {
-  assert(net);
-  assert(net->num_layers > 0);
-  return net->weights[0].rows;
+  return nnLayerInputSize(net, 0);
 }
 
 int nnNetOutputSize(const nnNeuralNetwork* net) {
-  assert(net);
-  assert(net->num_layers > 0);
-  return net->weights[net->num_layers - 1].cols;
+  return nnLayerOutputSize(net, net->num_layers - 1);
 }
 
-int nnLayerInputSize(const nnMatrix* weights) {
-  assert(weights);
-  return weights->rows;
+int nnLayerInputSize(const nnNeuralNetwork* net, int layer) {
+  assert(net);
+  return net->layers[layer].input_size;
 }
 
-int nnLayerOutputSize(const nnMatrix* weights) {
-  assert(weights);
-  return weights->cols;
+int nnLayerOutputSize(const nnNeuralNetwork* net, int layer) {
+  assert(net);
+  return net->layers[layer].output_size;
 }