xref: /aosp_15_r20/external/libopus/dnn/training_tf2/dump_plc.py (revision a58d3d2adb790c104798cd88c8a3aff4fa8b82cc)
1#!/usr/bin/python3
2'''Copyright (c) 2021-2022 Amazon
3   Copyright (c) 2017-2018 Mozilla
4
5   Redistribution and use in source and binary forms, with or without
6   modification, are permitted provided that the following conditions
7   are met:
8
9   - Redistributions of source code must retain the above copyright
10   notice, this list of conditions and the following disclaimer.
11
12   - Redistributions in binary form must reproduce the above copyright
13   notice, this list of conditions and the following disclaimer in the
14   documentation and/or other materials provided with the distribution.
15
16   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
17   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
18   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
19   A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR
20   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
21   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
22   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
23   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
24   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
25   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
26   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27'''
28
29import lpcnet_plc
30import io
31import sys
32import numpy as np
33from tensorflow.keras.optimizers import Adam
34from tensorflow.keras.layers import Layer, GRU, Dense, Conv1D, Embedding
35import h5py
36import re
37
38# Flag for dumping e2e (differentiable lpc) network weights
39flag_e2e = False
40
41max_rnn_neurons = 1
42max_conv_inputs = 1
43
44def printVector(f, vector, name, dtype='float', dotp=False):
45    global array_list
46    if dotp:
47        vector = vector.reshape((vector.shape[0]//4, 4, vector.shape[1]//8, 8))
48        vector = vector.transpose((2, 0, 3, 1))
49    v = np.reshape(vector, (-1));
50    #print('static const float ', name, '[', len(v), '] = \n', file=f)
51    if name not in array_list:
52        array_list.append(name)
53    f.write('#ifndef USE_WEIGHTS_FILE\n')
54    f.write('#define WEIGHTS_{}_DEFINED\n'.format(name))
55    f.write('#define WEIGHTS_{}_TYPE WEIGHT_TYPE_{}\n'.format(name, dtype))
56    f.write('static const {} {}[{}] = {{\n   '.format(dtype, name, len(v)))
57    for i in range(0, len(v)):
58        f.write('{}'.format(v[i]))
59        if (i!=len(v)-1):
60            f.write(',')
61        else:
62            break;
63        if (i%8==7):
64            f.write("\n   ")
65        else:
66            f.write(" ")
67    #print(v, file=f)
68    f.write('\n};\n')
69    f.write('#endif\n\n')
70    return;
71
72def printSparseVector(f, A, name, have_diag=True):
73    N = A.shape[0]
74    M = A.shape[1]
75    W = np.zeros((0,), dtype='int')
76    W0 = np.zeros((0,))
77    if have_diag:
78        diag = np.concatenate([np.diag(A[:,:N]), np.diag(A[:,N:2*N]), np.diag(A[:,2*N:])])
79        A[:,:N] = A[:,:N] - np.diag(np.diag(A[:,:N]))
80        A[:,N:2*N] = A[:,N:2*N] - np.diag(np.diag(A[:,N:2*N]))
81        A[:,2*N:] = A[:,2*N:] - np.diag(np.diag(A[:,2*N:]))
82        printVector(f, diag, name + '_diag')
83    AQ = np.minimum(127, np.maximum(-128, np.round(A*128))).astype('int')
84    idx = np.zeros((0,), dtype='int')
85    for i in range(M//8):
86        pos = idx.shape[0]
87        idx = np.append(idx, -1)
88        nb_nonzero = 0
89        for j in range(N//4):
90            block = A[j*4:(j+1)*4, i*8:(i+1)*8]
91            qblock = AQ[j*4:(j+1)*4, i*8:(i+1)*8]
92            if np.sum(np.abs(block)) > 1e-10:
93                nb_nonzero = nb_nonzero + 1
94                idx = np.append(idx, j*4)
95                vblock = qblock.transpose((1,0)).reshape((-1,))
96                W0 = np.concatenate([W0, block.reshape((-1,))])
97                W = np.concatenate([W, vblock])
98        idx[pos] = nb_nonzero
99    f.write('#ifdef DOT_PROD\n')
100    printVector(f, W, name, dtype='qweight')
101    f.write('#else /*DOT_PROD*/\n')
102    printVector(f, W0, name, dtype='qweight')
103    f.write('#endif /*DOT_PROD*/\n')
104    #idx = np.tile(np.concatenate([np.array([N]), np.arange(N)]), 3*N//16)
105    printVector(f, idx, name + '_idx', dtype='int')
106    return AQ
107
108def dump_layer_ignore(self, f, hf):
109    print("ignoring layer " + self.name + " of type " + self.__class__.__name__)
110    return False
111Layer.dump_layer = dump_layer_ignore
112
113def dump_sparse_gru(self, f, hf):
114    global max_rnn_neurons
115    name = 'sparse_' + self.name
116    print("printing layer " + name + " of type sparse " + self.__class__.__name__)
117    weights = self.get_weights()
118    qweights = printSparseVector(f, weights[1], name + '_recurrent_weights')
119    printVector(f, weights[-1], name + '_bias')
120    subias = weights[-1].copy()
121    subias[1,:] = subias[1,:] - np.sum(qweights*(1./128),axis=0)
122    printVector(f, subias, name + '_subias')
123    if hasattr(self, 'activation'):
124        activation = self.activation.__name__.upper()
125    else:
126        activation = 'TANH'
127    if hasattr(self, 'reset_after') and not self.reset_after:
128        reset_after = 0
129    else:
130        reset_after = 1
131    neurons = weights[0].shape[1]//3
132    max_rnn_neurons = max(max_rnn_neurons, neurons)
133    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
134    hf.write('#define {}_STATE_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
135    model_struct.write('  SparseGRULayer {};\n'.format(name));
136    model_init.write('  if (sparse_gru_init(&model->{}, arrays, "{}_bias", "{}_subias", "{}_recurrent_weights_diag", "{}_recurrent_weights", "{}_recurrent_weights_idx",  {}, ACTIVATION_{}, {})) return 1;\n'
137            .format(name, name, name, name, name, name, weights[0].shape[1]//3, activation, reset_after))
138    return True
139
140def dump_gru_layer(self, f, hf):
141    global max_rnn_neurons
142    name = self.name
143    print("printing layer " + name + " of type " + self.__class__.__name__)
144    weights = self.get_weights()
145    qweight = printSparseVector(f, weights[0], name + '_weights', have_diag=False)
146
147    f.write('#ifdef DOT_PROD\n')
148    qweight2 = np.clip(np.round(128.*weights[1]).astype('int'), -128, 127)
149    printVector(f, qweight2, name + '_recurrent_weights', dotp=True, dtype='qweight')
150    f.write('#else /*DOT_PROD*/\n')
151    printVector(f, weights[1], name + '_recurrent_weights')
152    f.write('#endif /*DOT_PROD*/\n')
153
154    printVector(f, weights[-1], name + '_bias')
155    subias = weights[-1].copy()
156    subias[0,:] = subias[0,:] - np.sum(qweight*(1./128.),axis=0)
157    subias[1,:] = subias[1,:] - np.sum(qweight2*(1./128.),axis=0)
158    printVector(f, subias, name + '_subias')
159    if hasattr(self, 'activation'):
160        activation = self.activation.__name__.upper()
161    else:
162        activation = 'TANH'
163    if hasattr(self, 'reset_after') and not self.reset_after:
164        reset_after = 0
165    else:
166        reset_after = 1
167    neurons = weights[0].shape[1]//3
168    max_rnn_neurons = max(max_rnn_neurons, neurons)
169    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
170    hf.write('#define {}_STATE_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
171    model_struct.write('  GRULayer {};\n'.format(name));
172    model_init.write('  if (gru_init(&model->{}, arrays, "{}_bias", "{}_subias", "{}_weights", "{}_weights_idx", "{}_recurrent_weights", {}, {}, ACTIVATION_{}, {})) return 1;\n'
173             .format(name, name, name, name, name, name, weights[0].shape[0], weights[0].shape[1]//3, activation, reset_after))
174    return True
175GRU.dump_layer = dump_gru_layer
176
177def dump_gru_layer_dummy(self, f, hf):
178    name = self.name
179    weights = self.get_weights()
180    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
181    hf.write('#define {}_STATE_SIZE {}\n'.format(name.upper(), weights[0].shape[1]//3))
182    return True;
183
184#GRU.dump_layer = dump_gru_layer_dummy
185
186def dump_dense_layer_impl(name, weights, bias, activation, f, hf):
187    printVector(f, weights, name + '_weights')
188    printVector(f, bias, name + '_bias')
189    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights.shape[1]))
190    model_struct.write('  DenseLayer {};\n'.format(name));
191    model_init.write('  if (dense_init(&model->{}, arrays, "{}_bias", "{}_weights", {}, {}, ACTIVATION_{})) return 1;\n'
192            .format(name, name, name, weights.shape[0], weights.shape[1], activation))
193
194def dump_dense_layer(self, f, hf):
195    name = self.name
196    print("printing layer " + name + " of type " + self.__class__.__name__)
197    weights = self.get_weights()
198    activation = self.activation.__name__.upper()
199    dump_dense_layer_impl(name, weights[0], weights[1], activation, f, hf)
200    return False
201
202Dense.dump_layer = dump_dense_layer
203
204def dump_conv1d_layer(self, f, hf):
205    global max_conv_inputs
206    name = self.name
207    print("printing layer " + name + " of type " + self.__class__.__name__)
208    weights = self.get_weights()
209    printVector(f, weights[0], name + '_weights')
210    printVector(f, weights[-1], name + '_bias')
211    activation = self.activation.__name__.upper()
212    max_conv_inputs = max(max_conv_inputs, weights[0].shape[1]*weights[0].shape[0])
213    hf.write('#define {}_OUT_SIZE {}\n'.format(name.upper(), weights[0].shape[2]))
214    hf.write('#define {}_STATE_SIZE ({}*{})\n'.format(name.upper(), weights[0].shape[1], (weights[0].shape[0]-1)))
215    hf.write('#define {}_DELAY {}\n'.format(name.upper(), (weights[0].shape[0]-1)//2))
216    model_struct.write('  Conv1DLayer {};\n'.format(name));
217    model_init.write('  if (conv1d_init(&model->{}, arrays, "{}_bias", "{}_weights", {}, {}, {}, ACTIVATION_{})) return 1;\n'
218            .format(name, name, name, weights[0].shape[1], weights[0].shape[0], weights[0].shape[2], activation))
219    return True
220Conv1D.dump_layer = dump_conv1d_layer
221
222
223
224filename = sys.argv[1]
225with h5py.File(filename, "r") as f:
226    units = min(f['model_weights']['plc_gru1']['plc_gru1']['recurrent_kernel:0'].shape)
227    units2 = min(f['model_weights']['plc_gru2']['plc_gru2']['recurrent_kernel:0'].shape)
228    cond_size = f['model_weights']['plc_dense1']['plc_dense1']['kernel:0'].shape[1]
229
230model = lpcnet_plc.new_lpcnet_plc_model(rnn_units=units, cond_size=cond_size)
231model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy'])
232#model.summary()
233
234model.load_weights(filename, by_name=True)
235
236if len(sys.argv) > 2:
237    cfile = sys.argv[2];
238    hfile = sys.argv[3];
239else:
240    cfile = 'plc_data.c'
241    hfile = 'plc_data.h'
242
243
244f = open(cfile, 'w')
245hf = open(hfile, 'w')
246model_struct = io.StringIO()
247model_init = io.StringIO()
248model_struct.write('typedef struct {\n')
249model_init.write('#ifndef DUMP_BINARY_WEIGHTS\n')
250model_init.write('int init_plc_model(PLCModel *model, const WeightArray *arrays) {\n')
251array_list = []
252
253
254f.write('/*This file is automatically generated from a Keras model*/\n')
255f.write('/*based on model {}*/\n\n'.format(sys.argv[1]))
256f.write('#ifdef HAVE_CONFIG_H\n#include "config.h"\n#endif\n\n#include "nnet.h"\n#include "{}"\n\n'.format(hfile))
257
258hf.write('/*This file is automatically generated from a Keras model*/\n\n')
259hf.write('#ifndef PLC_DATA_H\n#define PLC_DATA_H\n\n#include "nnet.h"\n\n')
260
261layer_list = []
262for i, layer in enumerate(model.layers):
263    if layer.dump_layer(f, hf):
264        layer_list.append(layer.name)
265
266#dump_sparse_gru(model.get_layer('gru_a'), f, hf)
267f.write('#ifndef USE_WEIGHTS_FILE\n')
268f.write('const WeightArray lpcnet_plc_arrays[] = {\n')
269for name in array_list:
270    f.write('#ifdef WEIGHTS_{}_DEFINED\n'.format(name))
271    f.write('  {{"{}", WEIGHTS_{}_TYPE, sizeof({}), {}}},\n'.format(name, name, name, name))
272    f.write('#endif\n')
273f.write('  {NULL, 0, 0, NULL}\n};\n')
274f.write('#endif\n')
275
276model_init.write('  return 0;\n}\n')
277model_init.write('#endif\n')
278f.write(model_init.getvalue())
279
280
281hf.write('#define PLC_MAX_RNN_NEURONS {}\n\n'.format(max_rnn_neurons))
282#hf.write('#define PLC_MAX_CONV_INPUTS {}\n\n'.format(max_conv_inputs))
283
284hf.write('typedef struct {\n')
285for i, name in enumerate(layer_list):
286    hf.write('  float {}_state[{}_STATE_SIZE];\n'.format(name, name.upper()))
287hf.write('} PLCNetState;\n\n')
288
289model_struct.write('} PLCModel;\n\n')
290hf.write(model_struct.getvalue())
291hf.write('int init_plc_model(PLCModel *model, const WeightArray *arrays);\n\n')
292
293hf.write('\n\n#endif\n')
294
295f.close()
296hf.close()
297