adapted dumping API

5 files changed, 58 insertions, 288 deletions

diff --git a/dnn/torch/weight-exchange/setup.py b/dnn/torch/weight-exchange/setup.py
index bf08db19..e590aad6 100644
--- a/dnn/torch/weight-exchange/setup.py
+++ b/dnn/torch/weight-exchange/setup.py
@@ -39,7 +39,7 @@ with open(os.path.join(lib_folder, 'requirements.txt'), 'r') as f:
 print(install_requires)
 
 setup(name='wexchange',
-      version='1.4',
+      version='1.5',
       author='Jan Buethe',
       author_email='jbuethe@amazon.de',
       description='Weight-exchange library between Pytorch and Tensorflow',
diff --git a/dnn/torch/weight-exchange/wexchange/c_export/__init__.py b/dnn/torch/weight-exchange/wexchange/c_export/__init__.py
index bb9a5eef..331c2e05 100644
--- a/dnn/torch/weight-exchange/wexchange/c_export/__init__.py
+++ b/dnn/torch/weight-exchange/wexchange/c_export/__init__.py
@@ -28,4 +28,4 @@ from .c_writer import CWriter
 */
 """
 
-from .common import print_gru_layer, print_dense_layer, print_conv1d_layer, print_vector, print_gru_layer2
\ No newline at end of file
+from .common import print_gru_layer, print_dense_layer, print_conv1d_layer, print_vector
\ No newline at end of file
diff --git a/dnn/torch/weight-exchange/wexchange/c_export/common.py b/dnn/torch/weight-exchange/wexchange/c_export/common.py
index a1be5203..c95cd88d 100644
--- a/dnn/torch/weight-exchange/wexchange/c_export/common.py
+++ b/dnn/torch/weight-exchange/wexchange/c_export/common.py
@@ -161,165 +161,6 @@ def print_sparse_weight(writer, A, name, scale=1/128, have_diag=True, quantize=F
     return Aq
 
 
-
-
-
-def print_sparse_vector(writer, A, name, have_diag=True, quantize=False):
-    N = A.shape[0]
-    M = A.shape[1]
-    W = np.zeros((0,), dtype='int')
-    W0 = np.zeros((0,))
-    if have_diag:
-        diag = np.concatenate([np.diag(A[:,:N]), np.diag(A[:,N:2*N]), np.diag(A[:,2*N:])])
-        A[:,:N] = A[:,:N] - np.diag(np.diag(A[:,:N]))
-        A[:,N:2*N] = A[:,N:2*N] - np.diag(np.diag(A[:,N:2*N]))
-        A[:,2*N:] = A[:,2*N:] - np.diag(np.diag(A[:,2*N:]))
-        print_vector(writer, diag, name + '_diag')
-    AQ = np.minimum(127, np.maximum(-128, np.round(A*128))).astype('int')
-    idx = np.zeros((0,), dtype='int')
-    for i in range(M//8):
-        pos = idx.shape[0]
-        idx = np.append(idx, -1)
-        nb_nonzero = 0
-        for j in range(N//4):
-            block = A[j*4:(j+1)*4, i*8:(i+1)*8]
-            qblock = AQ[j*4:(j+1)*4, i*8:(i+1)*8]
-            if np.sum(np.abs(block)) > 1e-10:
-                nb_nonzero = nb_nonzero + 1
-                idx = np.append(idx, j*4)
-                vblock = qblock.transpose((1,0)).reshape((-1,))
-                W0 = np.concatenate([W0, block.reshape((-1,))])
-                W = np.concatenate([W, vblock])
-        idx[pos] = nb_nonzero
-
-    if quantize: print_vector(writer, W, name + '_int8', reshape_8x4=False, dtype='opus_int8')
-    print_vector(writer, W0, name + '_float', reshape_8x4=False, dtype='float', debug_float=quantize)
-    print_vector(writer, idx, name + '_idx', reshape_8x4=False, dtype='int')
-
-def print_sparse_vector(writer, A, name, have_diag=True, quantize=False):
-    N = A.shape[0]
-    M = A.shape[1]
-    W = np.zeros((0,), dtype='int')
-    W0 = np.zeros((0,))
-    if have_diag:
-        diag = np.concatenate([np.diag(A[:,:N]), np.diag(A[:,N:2*N]), np.diag(A[:,2*N:])])
-        A[:,:N] = A[:,:N] - np.diag(np.diag(A[:,:N]))
-        A[:,N:2*N] = A[:,N:2*N] - np.diag(np.diag(A[:,N:2*N]))
-        A[:,2*N:] = A[:,2*N:] - np.diag(np.diag(A[:,2*N:]))
-        print_vector(writer, diag, name + '_diag')
-    AQ = np.minimum(127, np.maximum(-128, np.round(A*128))).astype('int')
-    idx = np.zeros((0,), dtype='int')
-    for i in range(M//8):
-        pos = idx.shape[0]
-        idx = np.append(idx, -1)
-        nb_nonzero = 0
-        for j in range(N//4):
-            block = A[j*4:(j+1)*4, i*8:(i+1)*8]
-            qblock = AQ[j*4:(j+1)*4, i*8:(i+1)*8]
-            if np.sum(np.abs(block)) > 1e-10:
-                nb_nonzero = nb_nonzero + 1
-                idx = np.append(idx, j*4)
-                vblock = qblock.transpose((1,0)).reshape((-1,))
-                W0 = np.concatenate([W0, block.reshape((-1,))])
-                W = np.concatenate([W, vblock])
-        idx[pos] = nb_nonzero
-
-    if quantize: print_vector(writer, W, name + '_int8', reshape_8x4=False, dtype='opus_int8')
-    print_vector(writer, W0, name + '_float', reshape_8x4=False, dtype='float', debug_float=quantize)
-    print_vector(writer, idx, name + '_idx', reshape_8x4=False, dtype='int')
-
-    return AQ
-
-def _check_activation(activation):
-    if not activation in {"TANH", "SIGMOID", "LINEAR", "SWISH", "RELU", "SOFTMAX"}:
-        raise ValueError(f"error: unknown activation {activation}")
-
-def print_dense_layer(writer : CWriter,
-                      name : str,
-                      weight : np.ndarray,
-                      bias : np.ndarray,
-                      activation: str,
-                      format : str = 'torch'):
-
-    _check_activation(activation)
-
-    if format == 'torch':
-        weight = weight.transpose()
-
-    print_vector(writer, weight, name + "_weights")
-    print_vector(writer, bias, name + "_bias")
-
-    writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {weight.shape[1]}\n")
-
-    if writer.enable_binary_blob:
-        init_call = f'linear_init(&model->{name}, arrays, "{name}_bias", "{name}_weights", {weight.shape[0]}, {weight.shape[1]}, ACTIVATION_{activation})'
-        writer.layer_dict[name] = ('DenseLayer', init_call)
-    else:
-        writer.source.write(
-f"""
-
-const DenseLayer {name} = {{
-   {name}_bias,
-   {name}_weights,
-   {weight.shape[0]},
-   {weight.shape[1]},
-   ACTIVATION_{activation}
-}};
-
-"""
-        )
-
-        writer.header.write(f"\nextern const DenseLayer {name};\n\n")
-
-
-
-
-
-def print_conv1d_layer(writer : CWriter,
-                       name : str,
-                       weight : np.ndarray,
-                       bias : np.ndarray,
-                       activation: str,
-                       format : str = 'torch'):
-
-    _check_activation(activation)
-
-    if format == "torch":
-        # convert to channels last
-        weight = np.transpose(weight, (2, 1, 0))
-
-    print_vector(writer, weight, name + "_weights")
-    print_vector(writer, bias, name + "_bias")
-
-    writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {weight.shape[2]}\n")
-    writer.header.write(f"\n#define {name.upper()}_STATE_SIZE ({weight.shape[1]} * ({weight.shape[0] - 1}))\n")
-    writer.header.write(f"\n#define {name.upper()}_DELAY {(weight.shape[0] - 1) // 2}\n") # CAVE: delay is not a property of the conv layer
-
-    if writer.enable_binary_blob:
-        init_call = f'conv1d_init(&model->{name}, arrays, "{name}_bias", "{name}_weights", {weight.shape[1]}, {weight.shape[0]}, {weight.shape[2]}, ACTIVATION_{activation})'
-        writer.layer_dict[name] = ('Conv1DLayer', init_call)
-    else:
-
-        writer.source.write(
-f"""
-
-const Conv1DLayer {name} = {{
-   {name}_bias,
-   {name}_weights,
-   {weight.shape[1]},
-   {weight.shape[0]},
-   {weight.shape[2]},
-   ACTIVATION_{activation}
-}};
-
-"""
-        )
-
-        writer.header.write(f"\nextern const Conv1DLayer {name};\n\n")
-
-    return weight.shape[0] * weight.shape[1]
-
-
 def qn(string):
     if string == "NULL": return string
     else: return '"' + string + '"'
@@ -395,34 +236,47 @@ def print_linear_layer(writer : CWriter,
     writer.layer_dict[name] = ('LinearLayer', init_call)
 
 
-def print_gru_layer2(writer : CWriter,
-                    name : str,
-                    weight : np.ndarray,
-                    recurrent_weight : np.ndarray,
-                    bias : np.ndarray,
-                    recurrent_bias : np.ndarray,
-                    format : str = 'torch',
-                    quantize : bool = False,
-                    input_sparse : bool = False,
-                    recurrent_sparse : bool = False,
-                    scale=1/128
-                    ):
+def print_dense_layer(writer : CWriter,
+                      name : str,
+                      weight : np.ndarray,
+                      bias : np.ndarray,
+                      scale=1/128,
+                      format : str = 'torch',
+                      sparse=False,
+                      diagonal=False,
+                      quantize=False):
+
+
+    if format == 'torch':
+        weight = weight.transpose()
+
+    print_linear_layer(writer, name, weight, bias, scale=scale, format=format, sparse=sparse, diagonal=diagonal, quantize=quantize)
+
+    writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {weight.shape[1]}\n")
+
+
+def print_conv1d_layer(writer : CWriter,
+                       name : str,
+                       weight : np.ndarray,
+                       bias : np.ndarray,
+                       scale=1/128,
+                       format : str = 'torch',
+                       quantize=False):
+
 
     if format == "torch":
-        # change gate ordering from rzn to zrn
+        # convert to channels last
+        weight = np.transpose(weight, (2, 1, 0))
 
-        N = weight.shape[0] // 3
-        for x in [weight, recurrent_weight, bias, recurrent_bias]:
-            tmp = x[0:N].copy()
-            x[0:N] = x[N:2*N]
-            x[N:2*N] = tmp
+    lin_weight = np.reshape(weight, (-1, weight.shape[-1]))
+    print_linear_layer(writer, name, lin_weight, bias, scale=scale, format=format, sparse=False, diagonal=False, quantize=quantize)
 
-    print_linear_layer(writer, name + "_input", weight, bias, scale=scale, format=format, sparse=input_sparse, quantize=quantize)
-    print_linear_layer(writer, name + "_recurrent", recurrent_weight, recurrent_bias, scale=scale, format=format, sparse=recurrent_sparse, diagonal=recurrent_sparse, quantize=quantize)
+    writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {weight.shape[2]}\n")
+    writer.header.write(f"\n#define {name.upper()}_STATE_SIZE ({weight.shape[1]} * ({weight.shape[0] - 1}))\n")
+    writer.header.write(f"\n#define {name.upper()}_DELAY {(weight.shape[0] - 1) // 2}\n") # CAVE: delay is not a property of the conv layer
+
+    return weight.shape[0] * weight.shape[1]
 
-    # wrapping it up
-    writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {N}\n")
-    writer.header.write(f"\n#define {name.upper()}_STATE_SIZE {N}\n")
 
 def print_gru_layer(writer : CWriter,
                     name : str,
@@ -430,17 +284,16 @@ def print_gru_layer(writer : CWriter,
                     recurrent_weight : np.ndarray,
                     bias : np.ndarray,
                     recurrent_bias : np.ndarray,
-                    activation: str,
                     format : str = 'torch',
-                    dotp : bool = False,
+                    quantize : bool = False,
                     input_sparse : bool = False,
-                    reset_after : int = 0
+                    recurrent_sparse : bool = False,
+                    scale=1/128,
+                    recurrent_scale=1/128
                     ):
 
-    _check_activation(activation)
-
     if format == "torch":
-        # transpose weight matrices and change gate order from rzn to zrn
+        # change gate ordering from rzn to zrn
 
         N = weight.shape[0] // 3
         for x in [weight, recurrent_weight, bias, recurrent_bias]:
@@ -448,82 +301,9 @@ def print_gru_layer(writer : CWriter,
             x[0:N] = x[N:2*N]
             x[N:2*N] = tmp
 
-        weight = weight.transpose()
-        recurrent_weight = recurrent_weight.transpose()
-
-
-    # input weights
-    if input_sparse:
-        qweight = print_sparse_vector(writer, weight, name + '_weights', have_diag=False)
-    else:
-        qweight = np.clip(np.round(128. * weight).astype('int'), -128, 127)
-
-        if dotp:
-            writer.source.write("#ifdef DOT_PROD\n")
-            print_vector(writer, qweight, name + '_weights', dtype='qweight', dotp=True)
-            writer.source.write("#else /*DOT_PROD*/\n")
-
-        print_vector(writer, weight, name + '_weights')
-
-        if dotp:
-             writer.source.write("#endif /*DOT_PROD*/\n")
-
-
-    # recurrent weights
-    recurrent_qweight = np.clip(np.round(128. * recurrent_weight).astype('int'), -128, 127)
-
-    if dotp:
-        writer.source.write("#ifdef DOT_PROD\n")
-        print_vector(writer, recurrent_qweight, name + '_recurrent_weights', dtype='qweight', dotp=True)
-        writer.source.write("#else /*DOT_PROD*/\n")
-
-    print_vector(writer, recurrent_weight, name + '_recurrent_weights')
-
-    if dotp:
-        writer.source.write("#endif /*DOT_PROD*/\n")
-
-
-    # corrected bias for unsigned int matrix multiplication
-    subias              = bias - np.sum(qweight / 128., axis=0)
-    recurrent_subias    = recurrent_bias - np.sum(recurrent_qweight / 128., axis=0)
-
-    print_vector(writer, np.concatenate((bias, recurrent_bias)), name + "_bias")
-    print_vector(writer, np.concatenate((subias, recurrent_subias)), name + "_subias")
-
+    print_linear_layer(writer, name + "_input", weight, bias, scale=scale, format=format, sparse=input_sparse, quantize=quantize)
+    print_linear_layer(writer, name + "_recurrent", recurrent_weight, recurrent_bias, scale=recurrent_scale, format=format, sparse=recurrent_sparse, diagonal=recurrent_sparse, quantize=quantize)
 
     # wrapping it up
     writer.header.write(f"\n#define {name.upper()}_OUT_SIZE {N}\n")
-    writer.header.write(f"\n#define {name.upper()}_STATE_SIZE {N}\n")
-
-    if writer.enable_binary_blob:
-        if input_sparse:
-            init_call = f'gru_init(&model->{name}, arrays, "{name}_bias", "{name}_subias", "{name}_weights", "{name + "_weights_idx"}", "{name}_recurrent_weights", {weight.shape[0]}, {weight.shape[1] // 3}, ACTIVATION_{activation}, {reset_after})'
-        else:
-            init_call = f'gru_init(&model->{name}, arrays, "{name}_bias", "{name}_subias", "{name}_weights", NULL, "{name}_recurrent_weights", {weight.shape[0]}, {weight.shape[1] // 3}, ACTIVATION_{activation}, {reset_after})'
-
-        writer.layer_dict[name] = ('GRULayer', init_call)
-
-    else:
-
-        writer.source.write(
-f"""
-
-const GRULayer {name} = {{
-   {name}_bias,
-   {name}_subias,
-   {name}_weights,
-   {name + "_weights_idx" if input_sparse else "NULL"},
-   {name}_recurrent_weights,
-   {weight.shape[0]},
-   {weight.shape[1] // 3},
-   ACTIVATION_{activation},
-   {reset_after}
-}};
-
-"""
-        )
-
-        writer.header.write(f"\nextern const GRULayer {name};\n")
-
-
-    return N
+    writer.header.write(f"\n#define {name.upper()}_STATE_SIZE {N}\n")
\ No newline at end of file
diff --git a/dnn/torch/weight-exchange/wexchange/tf/tf.py b/dnn/torch/weight-exchange/wexchange/tf/tf.py
index c8f9ed2f..bebbb55a 100644
--- a/dnn/torch/weight-exchange/wexchange/tf/tf.py
+++ b/dnn/torch/weight-exchange/wexchange/tf/tf.py
@@ -34,7 +34,7 @@ import numpy as np
 
 from wexchange.c_export import CWriter, print_gru_layer, print_dense_layer, print_conv1d_layer
 
-def dump_tf_gru_weights(where, gru, name=None, input_sparse=False, dotp=False):
+def dump_tf_gru_weights(where, gru, name='gru', input_sparse=False, recurrent_sparse=False, quantize=False, scale=1/128, recurrent_scale=1/128):
 
 
     assert gru.activation == tf.keras.activations.tanh
@@ -47,7 +47,7 @@ def dump_tf_gru_weights(where, gru, name=None, input_sparse=False, dotp=False):
     b_hh = gru.weights[2].numpy()[1].copy()
 
     if isinstance(where, CWriter):
-        return print_gru_layer(where, name, w_ih, w_hh, b_ih, b_hh, 'TANH', format='tf', reset_after=1, input_sparse=input_sparse, dotp=dotp)
+        return print_gru_layer(where, name, w_ih, w_hh, b_ih, b_hh, format='tf', input_sparse=input_sparse, recurrent_sparse=recurrent_sparse, quantize=quantize, scale=scale, recurrent_scale=recurrent_scale)
     else:
         os.makedirs(where, exist_ok=True)
 
@@ -87,7 +87,7 @@ def load_tf_gru_weights(path, gru):
     gru.weights[2].assign(tf.convert_to_tensor(np.vstack((b_ih, b_hh))))
 
 
-def dump_tf_dense_weights(where, dense, name=None):
+def dump_tf_dense_weights(where, dense, name='dense', scale=1/128, sparse=False, diagonal=False, quantize=False):
 
     w = dense.weights[0].numpy()
     if dense.bias is None:
@@ -98,12 +98,7 @@ def dump_tf_dense_weights(where, dense, name=None):
 
 
     if isinstance(where, CWriter):
-        try:
-            activation = dense.activation.__name__.upper()
-        except:
-            activation = "LINEAR"
-
-        return print_dense_layer(where, name, w, b, activation, format='tf')
+        return print_dense_layer(where, name, w, b, scale=scale, format='tf', sparse=sparse, diagonal=diagonal, quantize=quantize)
 
     else:
         os.makedirs(where, exist_ok=True)
@@ -122,7 +117,7 @@ def load_tf_dense_weights(path, dense):
         dense.weights[1].assign(tf.convert_to_tensor(b))
 
 
-def dump_tf_conv1d_weights(where, conv, name=None):
+def dump_tf_conv1d_weights(where, conv, name='conv', scale=1/128, quantize=False):
 
     assert conv.data_format == 'channels_last'
 
@@ -133,12 +128,7 @@ def dump_tf_conv1d_weights(where, conv, name=None):
         b = conv.bias.numpy()
 
     if isinstance(where, CWriter):
-        try:
-            activation = conv.activation.__name__.upper()
-        except:
-            activation = "LINEAR"
-
-        return print_conv1d_layer(where, name, w, b, activation, format='tf')
+        return print_conv1d_layer(where, name, w, b, scale=scale, format='tf', quantize=quantize)
     else:
         os.makedirs(where, exist_ok=True)
 
diff --git a/dnn/torch/weight-exchange/wexchange/torch/torch.py b/dnn/torch/weight-exchange/wexchange/torch/torch.py
index 60eef190..65c58d92 100644
--- a/dnn/torch/weight-exchange/wexchange/torch/torch.py
+++ b/dnn/torch/weight-exchange/wexchange/torch/torch.py
@@ -32,7 +32,7 @@ import os
 import torch
 import numpy as np
 
-from wexchange.c_export import CWriter, print_gru_layer2, print_dense_layer, print_conv1d_layer
+from wexchange.c_export import CWriter, print_gru_layer, print_dense_layer, print_conv1d_layer
 
 def dump_torch_gru_weights(where, gru, name='gru', input_sparse=False, recurrent_sparse=False, quantize=False, scale=1/128, recurrent_scale=1/128):
 
@@ -45,7 +45,7 @@ def dump_torch_gru_weights(where, gru, name='gru', input_sparse=False, recurrent
     b_hh = gru.bias_hh_l0.detach().cpu().numpy()
 
     if isinstance(where, CWriter):
-        return print_gru_layer2(where, name, w_ih, w_hh, b_ih, b_hh, format='torch', input_sparse=input_sparse, recurrent_sparse=recurrent_sparse, quantize=quantize, scale=scale)
+        return print_gru_layer(where, name, w_ih, w_hh, b_ih, b_hh, format='torch', input_sparse=input_sparse, recurrent_sparse=recurrent_sparse, quantize=quantize, scale=scale, recurrent_scale=recurrent_scale)
     else:
         os.makedirs(where, exist_ok=True)
 
@@ -73,7 +73,7 @@ def load_torch_gru_weights(where, gru):
         gru.bias_hh_l0.set_(torch.from_numpy(b_hh))
 
 
-def dump_torch_dense_weights(where, dense, name=None, activation="LINEAR"):
+def dump_torch_dense_weights(where, dense, name='dense', scale=1/128, sparse=False, diagonal=False, quantize=False):
 
     w = dense.weight.detach().cpu().numpy()
     if dense.bias is None:
@@ -82,7 +82,7 @@ def dump_torch_dense_weights(where, dense, name=None, activation="LINEAR"):
         b = dense.bias.detach().cpu().numpy()
 
     if isinstance(where, CWriter):
-        return print_dense_layer(where, name, w, b, activation, format='torch')
+        return print_dense_layer(where, name, w, b, scale=scale, format='torch', sparse=sparse, diagonal=diagonal, quantize=quantize)
 
     else:
         os.makedirs(where, exist_ok=True)
@@ -102,7 +102,7 @@ def load_torch_dense_weights(where, dense):
             dense.bias.set_(torch.from_numpy(b))
 
 
-def dump_torch_conv1d_weights(where, conv, name=None, activation="LINEAR"):
+def dump_torch_conv1d_weights(where, conv, name='conv', scale=1/128, quantize=False):
 
     w = conv.weight.detach().cpu().numpy()
     if conv.bias is None:
@@ -112,7 +112,7 @@ def dump_torch_conv1d_weights(where, conv, name=None, activation="LINEAR"):
 
     if isinstance(where, CWriter):
 
-        return print_conv1d_layer(where, name, w, b, activation, format='torch')
+        return print_conv1d_layer(where, name, w, b, scale=scale, format='torch', quantize=quantize)
     else:
         os.makedirs(where, exist_ok=True)