xref: /aosp_15_r20/external/libopus/dnn/training_tf2/dump_lpcnet.py (revision a58d3d2adb790c104798cd88c8a3aff4fa8b82cc)
1#!/usr/bin/python3
2'''Copyright (c) 2017-2018 Mozilla
3
4   Redistribution and use in source and binary forms, with or without
5   modification, are permitted provided that the following conditions
6   are met:
7
8   - Redistributions of source code must retain the above copyright
9   notice, this list of conditions and the following disclaimer.
10
11   - Redistributions in binary form must reproduce the above copyright
12   notice, this list of conditions and the following disclaimer in the
13   documentation and/or other materials provided with the distribution.
14
15   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
16   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
17   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
18   A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR
19   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
23   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
24   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26'''
27
28import os
29import io
30import lpcnet
31import sys
32import numpy as np
33from tensorflow.keras.optimizers import Adam
34from tensorflow.keras.layers import Layer, GRU, Dense, Conv1D, Embedding
35from ulaw import ulaw2lin, lin2ulaw
36from mdense import MDense
37from diffembed import diff_Embed
38from parameters import get_parameter
39import h5py
40import re
41import argparse
42
43
44# no cuda devices needed
45os.environ['CUDA_VISIBLE_DEVICES'] = ""
46
47# Flag for dumping e2e (differentiable lpc) network weights
48flag_e2e = False
49
50
51max_rnn_neurons = 1
52max_conv_inputs = 1
53max_mdense_tmp = 1
54
55def printVector(f, vector, name, dtype='float', dotp=False):
56    global array_list
57    if dotp:
58        vector = vector.reshape((vector.shape[0]//4, 4, vector.shape[1]//8, 8))
59        vector = vector.transpose((2, 0, 3, 1))
60    v = np.reshape(vector, (-1));
61    #print('static const float ', name, '[', len(v), '] = \n', file=f)
62    if name not in array_list:
63        array_list.append(name)
64    f.write('#ifndef USE_WEIGHTS_FILE\n')
65    f.write('#define WEIGHTS_{}_DEFINED\n'.format(name))
66    f.write('#define WEIGHTS_{}_TYPE WEIGHT_TYPE_{}\n'.format(name, dtype))
67    f.write('static const {} {}[{}] = {{\n   '.format(dtype, name, len(v)))
68    for i in range(0, len(v)):
69        f.write('{}'.format(v[i]))
70        if (i!=len(v)-1):
71            f.write(',')
72        else:
73            break;
74        if (i%8==7):
75            f.write("\n   ")
76        else:
77            f.write(" ")
78    #print(v, file=f)
79    f.write('\n};\n')
80    f.write('#endif\n\n')
81    return;
82
83def printSparseVector(f, A, name, have_diag=True):
84    N = A.shape[0]
85    M = A.shape[1]
86    W = np.zeros((0,), dtype='int')
87    W0 = np.zeros((0,))
88    if have_diag:
89        diag = np.concatenate([np.diag(A[:,:N]), np.diag(A[:,N:2*N]), np.diag(A[:,2*N:])])
90        A[:,:N] = A[:,:N] - np.diag(np.diag(A[:,:N]))
91        A[:,N:2*N] = A[:,N:2*N] - np.diag(np.diag(A[:,N:2*N]))
92        A[:,2*N:] = A[:,2*N:] - np.diag(np.diag(A[:,2*N:]))
93        printVector(f, diag, name + '_diag')
94    AQ = np.minimum(127, np.maximum(-128, np.round(A*128))).astype('int')
95    idx = np.zeros((0,), dtype='int')
96    for i in range(M//8):
97        pos = idx.shape[0]
98        idx = np.append(idx, -1)
99        nb_nonzero = 0
100        for j in range(N//4):
101            block = A[j*4:(j+1)*4, i*8:(i+1)*8]
102            qblock = AQ[j*4:(j+1)*4, i*8:(i+1)*8]
103            if np.sum(np.abs(block)) > 1e-10:
104                nb_nonzero = nb_nonzero + 1
105                idx = np.append(idx, j*4)
106                vblock = qblock.transpose((1,0)).reshape((-1,))
107                W0 = np.concatenate([W0, block.reshape((-1,))])
108                W = np.concatenate([W, vblock])
109        idx[pos] = nb_nonzero
110    f.write('#ifdef DOT_PROD\n')
111    printVector(f, W, name, dtype='qweight')
112    f.write('#else /*DOT_PROD*/\n')
113    printVector(f, W0, name, dtype='qweight')
114    f.write('#endif /*DOT_PROD*/\n')
115    #idx = np.tile(np.concatenate([np.array([N]), np.arange(N)]), 3*N//16)
116    printVector(f, idx, name + '_idx', dtype='int')
117    return AQ
118
119def dump_layer_ignore(self, f, hf):
120    print("ignoring layer " + self.name + " of type " + self.__class__.__name__)
121    return False
122Layer.dump_layer = dump_layer_ignore
123
124def dump_sparse_gru(self, f, hf):
125    global max_rnn_neurons
126    name = 'sparse_' + self.name
127    print("printing layer " + name + " of type sparse " + self.__class__.__name__)
128    weights = self.get_weights()
129    qweights = printSparseVector(f, weights[1], name + '_recurrent_weights')
130    printVector(f, weights[-1], name + '_bias')
131    subias = weights[-1].copy()
132    subias[1,:] = subias[1,:] - np.sum(qweights*(1./128),axis=0)
133    printVector(f, subias, name + '_subias')
134    if hasattr(self, 'activation'):
135        activation = self.activation.__name__.upper()
136    else:
137        activation = 'TANH'
138    if hasattr(self, 'reset_after') and not self.reset_after:
139        reset_after = 0
140    else:
141        reset_after = 1
142    neurons = weights[0].shape[1]//3
143    max_rnn_neurons = max(max_rnn_neurons, neurons)
144    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
145    hf.write('#define {}_STATE_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
146    model_struct.write('  SparseGRULayer {};\n'.format(name));
147    model_init.write('  if (sparse_gru_init(&model->{}, arrays, "{}_bias", "{}_subias", "{}_recurrent_weights_diag", "{}_recurrent_weights", "{}_recurrent_weights_idx",  {}, ACTIVATION_{}, {})) return 1;\n'
148            .format(name, name, name, name, name, name, weights[0].shape[1]//3, activation, reset_after))
149    return True
150
151def dump_grub(self, f, hf, gru_a_size):
152    global max_rnn_neurons
153    name = self.name
154    print("printing layer " + name + " of type " + self.__class__.__name__)
155    weights = self.get_weights()
156    qweight = printSparseVector(f, weights[0][:gru_a_size, :], name + '_weights', have_diag=False)
157
158    f.write('#ifdef DOT_PROD\n')
159    qweight2 = np.clip(np.round(128.*weights[1]).astype('int'), -128, 127)
160    printVector(f, qweight2, name + '_recurrent_weights', dotp=True, dtype='qweight')
161    f.write('#else /*DOT_PROD*/\n')
162    printVector(f, weights[1], name + '_recurrent_weights')
163    f.write('#endif /*DOT_PROD*/\n')
164
165    printVector(f, weights[-1], name + '_bias')
166    subias = weights[-1].copy()
167    subias[0,:] = subias[0,:] - np.sum(qweight*(1./128.),axis=0)
168    subias[1,:] = subias[1,:] - np.sum(qweight2*(1./128.),axis=0)
169    printVector(f, subias, name + '_subias')
170    if hasattr(self, 'activation'):
171        activation = self.activation.__name__.upper()
172    else:
173        activation = 'TANH'
174    if hasattr(self, 'reset_after') and not self.reset_after:
175        reset_after = 0
176    else:
177        reset_after = 1
178    neurons = weights[0].shape[1]//3
179    max_rnn_neurons = max(max_rnn_neurons, neurons)
180    model_struct.write('  GRULayer {};\n'.format(name));
181    model_init.write('  if (gru_init(&model->{}, arrays, "{}_bias", "{}_subias", "{}_weights", "{}_weights_idx", "{}_recurrent_weights", {}, {}, ACTIVATION_{}, {})) return 1;\n'
182            .format(name, name, name, name, name, name, gru_a_size, weights[0].shape[1]//3, activation, reset_after))
183    return True
184
185def dump_gru_layer_dummy(self, f, hf):
186    name = self.name
187    weights = self.get_weights()
188    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
189    hf.write('#define {}_STATE_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
190    return True;
191
192GRU.dump_layer = dump_gru_layer_dummy
193
194def dump_dense_layer_impl(name, weights, bias, activation, f, hf):
195    printVector(f, weights, name + '_weights')
196    printVector(f, bias, name + '_bias')
197    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights.shape[1]))
198    model_struct.write('  DenseLayer {};\n'.format(name));
199    model_init.write('  if (dense_init(&model->{}, arrays, "{}_bias", "{}_weights", {}, {}, ACTIVATION_{})) return 1;\n'
200            .format(name, name, name, weights.shape[0], weights.shape[1], activation))
201
202def dump_dense_layer(self, f, hf):
203    name = self.name
204    print("printing layer " + name + " of type " + self.__class__.__name__)
205    weights = self.get_weights()
206    activation = self.activation.__name__.upper()
207    dump_dense_layer_impl(name, weights[0], weights[1], activation, f, hf)
208    return False
209
210Dense.dump_layer = dump_dense_layer
211
212def dump_mdense_layer(self, f, hf):
213    global max_mdense_tmp
214    name = self.name
215    print("printing layer " + name + " of type " + self.__class__.__name__)
216    weights = self.get_weights()
217    printVector(f, np.transpose(weights[0], (0, 2, 1)), name + '_weights')
218    printVector(f, np.transpose(weights[1], (1, 0)), name + '_bias')
219    printVector(f, np.transpose(weights[2], (1, 0)), name + '_factor')
220    activation = self.activation.__name__.upper()
221    max_mdense_tmp = max(max_mdense_tmp, weights[0].shape[0]*weights[0].shape[2])
222    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[0]))
223    model_struct.write('  MDenseLayer {};\n'.format(name));
224    model_init.write('  if (mdense_init(&model->{}, arrays, "{}_bias",  "{}_weights",  "{}_factor",  {}, {}, {}, ACTIVATION_{})) return 1;\n'
225            .format(name, name, name, name, weights[0].shape[1], weights[0].shape[0], weights[0].shape[2], activation))
226    return False
227MDense.dump_layer = dump_mdense_layer
228
229def dump_conv1d_layer(self, f, hf):
230    global max_conv_inputs
231    name = self.name
232    print("printing layer " + name + " of type " + self.__class__.__name__)
233    weights = self.get_weights()
234    printVector(f, weights[0], name + '_weights')
235    printVector(f, weights[-1], name + '_bias')
236    activation = self.activation.__name__.upper()
237    max_conv_inputs = max(max_conv_inputs, weights[0].shape[1]*weights[0].shape[0])
238    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[2]))
239    hf.write('#define {}_STATE_SIZE ({}*{})\n'.format(name.upper(), weights[0].shape[1], (weights[0].shape[0]-1)))
240    hf.write('#define {}_DELAY {}\n'.format(name.upper(), (weights[0].shape[0]-1)//2))
241    model_struct.write('  Conv1DLayer {};\n'.format(name));
242    model_init.write('  if (conv1d_init(&model->{}, arrays, "{}_bias", "{}_weights", {}, {}, {}, ACTIVATION_{})) return 1;\n'
243            .format(name, name, name, weights[0].shape[1], weights[0].shape[0], weights[0].shape[2], activation))
244    return True
245Conv1D.dump_layer = dump_conv1d_layer
246
247
248def dump_embedding_layer_impl(name, weights, f, hf):
249    printVector(f, weights, name + '_weights')
250    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights.shape[1]))
251    model_struct.write('  EmbeddingLayer {};\n'.format(name));
252    model_init.write('  if (embedding_init(&model->{}, arrays, "{}_weights", {}, {})) return 1;\n'
253            .format(name, name, weights.shape[0], weights.shape[1]))
254
255def dump_embedding_layer(self, f, hf):
256    name = self.name
257    print("printing layer " + name + " of type " + self.__class__.__name__)
258    weights = self.get_weights()[0]
259    dump_embedding_layer_impl(name, weights, f, hf)
260    return False
261Embedding.dump_layer = dump_embedding_layer
262diff_Embed.dump_layer = dump_embedding_layer
263
264if __name__ == "__main__":
265    parser = argparse.ArgumentParser()
266    parser.add_argument('model_file', type=str, help='model weight h5 file')
267    parser.add_argument('--nnet-header', type=str, help='name of c header file for dumped model', default='nnet_data.h')
268    parser.add_argument('--nnet-source', type=str, help='name of c source file for dumped model', default='nnet_data.c')
269    parser.add_argument('--lpc-gamma', type=float, help='LPC weighting factor. If not specified I will attempt to read it from the model file with 1 as default', default=None)
270    parser.add_argument('--lookahead', type=float, help='Features lookahead. If not specified I will attempt to read it from the model file with 2 as default', default=None)
271
272    args = parser.parse_args()
273
274    filename = args.model_file
275    with h5py.File(filename, "r") as f:
276        units = min(f['model_weights']['gru_a']['gru_a']['recurrent_kernel:0'].shape)
277        units2 = min(f['model_weights']['gru_b']['gru_b']['recurrent_kernel:0'].shape)
278        cond_size = min(f['model_weights']['feature_dense1']['feature_dense1']['kernel:0'].shape)
279        e2e = 'rc2lpc' in f['model_weights']
280
281    model, _, _ = lpcnet.new_lpcnet_model(rnn_units1=units, rnn_units2=units2, flag_e2e = e2e, cond_size=cond_size)
282    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'])
283    #model.summary()
284
285    model.load_weights(filename, by_name=True)
286
287    cfile = args.nnet_source
288    hfile = args.nnet_header
289
290    f = open(cfile, 'w')
291    hf = open(hfile, 'w')
292    model_struct = io.StringIO()
293    model_init = io.StringIO()
294    model_struct.write('typedef struct {\n')
295    model_init.write('#ifndef DUMP_BINARY_WEIGHTS\n')
296    model_init.write('int init_lpcnet_model(LPCNetModel *model, const WeightArray *arrays) {\n')
297    array_list = []
298
299    f.write('/*This file is automatically generated from a Keras model*/\n')
300    f.write('/*based on model {}*/\n\n'.format(sys.argv[1]))
301    f.write('#ifdef HAVE_CONFIG_H\n#include "config.h"\n#endif\n\n#include "nnet.h"\n#include "{}"\n\n'.format(hfile))
302
303    hf.write('/*This file is automatically generated from a Keras model*/\n\n')
304    hf.write('#ifndef RNN_DATA_H\n#define RNN_DATA_H\n\n#include "nnet.h"\n\n')
305
306    if e2e:
307        hf.write('/* This is an end-to-end model */\n')
308        hf.write('#define END2END\n\n')
309    else:
310        hf.write('/* This is *not* an end-to-end model */\n')
311        hf.write('/* #define END2END */\n\n')
312
313    # LPC weighting factor
314    if type(args.lpc_gamma) == type(None):
315        lpc_gamma = get_parameter(model, 'lpc_gamma', 1)
316    else:
317        lpc_gamma = args.lpc_gamma
318
319    hf.write('/* LPC weighting factor */\n')
320    hf.write('#define LPC_GAMMA ' + str(lpc_gamma) +'f\n\n')
321
322    # look-ahead
323    if type(args.lookahead) == type(None):
324        lookahead = get_parameter(model, 'lookahead', 2)
325    else:
326        lookahead = args.lookahead
327
328    hf.write('/* Features look-ahead */\n')
329    hf.write('#define FEATURES_DELAY ' + str(lookahead) +'\n\n')
330
331    embed_size = lpcnet.embed_size
332
333    E = model.get_layer('embed_sig').get_weights()[0]
334    W = model.get_layer('gru_a').get_weights()[0][:embed_size,:]
335    dump_embedding_layer_impl('gru_a_embed_sig', np.dot(E, W), f, hf)
336    W = model.get_layer('gru_a').get_weights()[0][embed_size:2*embed_size,:]
337    dump_embedding_layer_impl('gru_a_embed_pred', np.dot(E, W), f, hf)
338    W = model.get_layer('gru_a').get_weights()[0][2*embed_size:3*embed_size,:]
339    dump_embedding_layer_impl('gru_a_embed_exc', np.dot(E, W), f, hf)
340    W = model.get_layer('gru_a').get_weights()[0][3*embed_size:,:]
341    #FIXME: dump only half the biases
342    b = model.get_layer('gru_a').get_weights()[2]
343    dump_dense_layer_impl('gru_a_dense_feature', W, b[:len(b)//2], 'LINEAR', f, hf)
344
345    W = model.get_layer('gru_b').get_weights()[0][model.rnn_units1:,:]
346    b = model.get_layer('gru_b').get_weights()[2]
347    # Set biases to zero because they'll be included in the GRU input part
348    # (we need regular and SU biases)
349    dump_dense_layer_impl('gru_b_dense_feature', W, 0*b[:len(b)//2], 'LINEAR', f, hf)
350    dump_grub(model.get_layer('gru_b'), f, hf, model.rnn_units1)
351
352    layer_list = []
353    for i, layer in enumerate(model.layers):
354        if layer.dump_layer(f, hf):
355            layer_list.append(layer.name)
356
357    dump_sparse_gru(model.get_layer('gru_a'), f, hf)
358
359    f.write('#ifndef USE_WEIGHTS_FILE\n')
360    f.write('const WeightArray lpcnet_arrays[] = {\n')
361    for name in array_list:
362        f.write('#ifdef WEIGHTS_{}_DEFINED\n'.format(name))
363        f.write('  {{"{}", WEIGHTS_{}_TYPE, sizeof({}), {}}},\n'.format(name, name, name, name))
364        f.write('#endif\n')
365    f.write('  {NULL, 0, 0, NULL}\n};\n')
366    f.write('#endif\n')
367
368    model_init.write('  return 0;\n}\n')
369    model_init.write('#endif\n')
370    f.write(model_init.getvalue())
371
372    hf.write('#define MAX_RNN_NEURONS {}\n\n'.format(max_rnn_neurons))
373    hf.write('#define MAX_CONV_INPUTS {}\n\n'.format(max_conv_inputs))
374    hf.write('#define MAX_MDENSE_TMP {}\n\n'.format(max_mdense_tmp))
375
376
377    hf.write('typedef struct {\n')
378    for i, name in enumerate(layer_list):
379        hf.write('  float {}_state[{}_STATE_SIZE];\n'.format(name, name.upper()))
380    hf.write('} NNetState;\n\n')
381
382    model_struct.write('} LPCNetModel;\n\n')
383    hf.write(model_struct.getvalue())
384    hf.write('int init_lpcnet_model(LPCNetModel *model, const WeightArray *arrays);\n\n')
385    hf.write('\n\n#endif\n')
386
387    f.close()
388    hf.close()
389