Some fixes to make model files work.

Author: GregorR

src/rnn_reader.c

@@ -36,6 +36,12 @@
 #include "rnn_data.h"
 #include "rnnoise.h"
 
+/* Although these values are the same as in rnn.h, we make them separate to
+ * avoid accidentally burning internal values into a file format */
+#define F_ACTIVATION_TANH       0
+#define F_ACTIVATION_SIGMOID    1
+#define F_ACTIVATION_RELU       2
+
 RNNModel *rnnoise_model_from_file(FILE *f)
 {
     int i, in;
@@ -71,6 +77,21 @@ RNNModel *rnnoise_model_from_file(FILE *f)
     name = in; \
     } while (0)
 
+#define INPUT_ACTIVATION(name) do { \
+    int activation; \
+    INPUT_VAL(activation); \
+    switch (activation) { \
+        case F_ACTIVATION_SIGMOID: \
+            name = ACTIVATION_SIGMOID; \
+            break; \
+        case F_ACTIVATION_RELU: \
+            name = ACTIVATION_RELU; \
+            break; \
+        default: \
+            name = ACTIVATION_TANH; \
+    } \
+    } while (0)
+
 #define INPUT_ARRAY(name, len) do { \
     rnn_weight *values = malloc((len) * sizeof(rnn_weight)); \
     if (!values) { \
@@ -90,13 +111,15 @@ RNNModel *rnnoise_model_from_file(FILE *f)
 #define INPUT_DENSE(name) do { \
     INPUT_VAL(name->nb_inputs); \
     INPUT_VAL(name->nb_neurons); \
+    INPUT_ACTIVATION(name->activation); \
     INPUT_ARRAY(name->input_weights, name->nb_inputs * name->nb_neurons); \
     INPUT_ARRAY(name->bias, name->nb_neurons); \
     } while (0)
 
 #define INPUT_GRU(name) do { \
     INPUT_VAL(name->nb_inputs); \
     INPUT_VAL(name->nb_neurons); \
+    INPUT_ACTIVATION(name->activation); \
     INPUT_ARRAY(name->input_weights, name->nb_inputs * name->nb_neurons * 3); \
     INPUT_ARRAY(name->recurrent_weights, name->nb_neurons * name->nb_neurons * 3); \
     INPUT_ARRAY(name->bias, name->nb_neurons * 3); \

training/dump_rnn.py

@@ -35,16 +35,22 @@ def printVector(f, ft, vector, name):
 
 def printLayer(f, ft, layer):
     weights = layer.get_weights()
+    activation = re.search('function (.*) at', str(layer.activation)).group(1).upper()
     if len(weights) > 2:
-        ft.write('{} {}\n'.format(weights[0].shape[0], weights[0].shape[1]/3))
+        ft.write('{} {} '.format(weights[0].shape[0], weights[0].shape[1]/3))
+    else:
+        ft.write('{} {} '.format(weights[0].shape[0], weights[0].shape[1]))
+    if activation == 'SIGMOID':
+        ft.write('1\n')
+    elif activation == 'RELU':
+        ft.write('2\n')
     else:
-        ft.write('{} {}\n'.format(weights[0].shape[0], weights[0].shape[1]))
+        ft.write('0\n')
     printVector(f, ft, weights[0], layer.name + '_weights')
     if len(weights) > 2:
         printVector(f, ft, weights[1], layer.name + '_recurrent_weights')
     printVector(f, ft, weights[-1], layer.name + '_bias')
     name = layer.name
-    activation = re.search('function (.*) at', str(layer.activation)).group(1).upper()
     if len(weights) > 2:
         f.write('static const GRULayer {} = {{\n   {}_bias,\n   {}_weights,\n   {}_recurrent_weights,\n   {}, {}, ACTIVATION_{}\n}};\n\n'
                 .format(name, name, name, name, weights[0].shape[0], weights[0].shape[1]/3, activation))
@@ -77,7 +83,7 @@ def mean_squared_sqrt_error(y_true, y_pred):
 ft = open(sys.argv[3], 'w')
 
 f.write('/*This file is automatically generated from a Keras model*/\n\n')
-f.write('#ifdef HAVE_CONFIG_H\n#include "config.h"\n#endif\n\n#include "rnn.h"\n\n')
+f.write('#ifdef HAVE_CONFIG_H\n#include "config.h"\n#endif\n\n#include "rnn.h"\n#include "rnn_data.h"\n\n')
 ft.write('rnnoise-nu model file version 1\n')
 
 layer_list = []