You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.6 KiB
108 lines
3.6 KiB
#!/usr/bin/python
|
|
|
|
from __future__ import print_function
|
|
|
|
from keras.models import Sequential
|
|
from keras.layers import Dense
|
|
from keras.layers import LSTM
|
|
from keras.layers import GRU
|
|
from keras.models import load_model
|
|
from keras import backend as K
|
|
import sys
|
|
import re
|
|
import numpy as np
|
|
|
|
def printVector(f, ft, vector, name):
|
|
v = np.reshape(vector, (-1));
|
|
#print('static const float ', name, '[', len(v), '] = \n', file=f)
|
|
f.write('static const rnn_weight {}[{}] = {{\n '.format(name, len(v)))
|
|
for i in range(0, len(v)):
|
|
f.write('{}'.format(min(127, int(round(256*v[i])))))
|
|
ft.write('{}'.format(min(127, int(round(256*v[i])))))
|
|
if (i!=len(v)-1):
|
|
f.write(',')
|
|
else:
|
|
break;
|
|
ft.write(" ")
|
|
if (i%8==7):
|
|
f.write("\n ")
|
|
else:
|
|
f.write(" ")
|
|
#print(v, file=f)
|
|
f.write('\n};\n\n')
|
|
ft.write("\n")
|
|
return;
|
|
|
|
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('{} {} '.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('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
|
|
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))
|
|
else:
|
|
f.write('static const DenseLayer {} = {{\n {}_bias,\n {}_weights,\n {}, {}, ACTIVATION_{}\n}};\n\n'
|
|
.format(name, name, name, weights[0].shape[0], weights[0].shape[1], activation))
|
|
|
|
def structLayer(f, layer):
|
|
weights = layer.get_weights()
|
|
name = layer.name
|
|
if len(weights) > 2:
|
|
f.write(' {},\n'.format(weights[0].shape[1]/3))
|
|
else:
|
|
f.write(' {},\n'.format(weights[0].shape[1]))
|
|
f.write(' &{},\n'.format(name))
|
|
|
|
|
|
def foo(c, name):
|
|
return None
|
|
|
|
def mean_squared_sqrt_error(y_true, y_pred):
|
|
return K.mean(K.square(K.sqrt(y_pred) - K.sqrt(y_true)), axis=-1)
|
|
|
|
|
|
model = load_model(sys.argv[1], custom_objects={'msse': mean_squared_sqrt_error, 'mean_squared_sqrt_error': mean_squared_sqrt_error, 'my_crossentropy': mean_squared_sqrt_error, 'mycost': mean_squared_sqrt_error, 'WeightClip': foo})
|
|
|
|
weights = model.get_weights()
|
|
|
|
f = open(sys.argv[2], 'w')
|
|
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#include "rnn_data.h"\n\n')
|
|
ft.write('rnnoise-nu model file version 1\n')
|
|
|
|
layer_list = []
|
|
for i, layer in enumerate(model.layers):
|
|
if len(layer.get_weights()) > 0:
|
|
printLayer(f, ft, layer)
|
|
if len(layer.get_weights()) > 2:
|
|
layer_list.append(layer.name)
|
|
|
|
f.write('const struct RNNModel rnnoise_model_{} = {{\n'.format(sys.argv[4]))
|
|
for i, layer in enumerate(model.layers):
|
|
if len(layer.get_weights()) > 0:
|
|
structLayer(f, layer)
|
|
f.write('};\n')
|
|
|
|
#hf.write('struct RNNState {\n')
|
|
#for i, name in enumerate(layer_list):
|
|
# hf.write(' float {}_state[{}_SIZE];\n'.format(name, name.upper()))
|
|
#hf.write('};\n')
|
|
|
|
f.close()
|