From a0e9e098c7e40a927294b8419cc73d11fe05f4ad Mon Sep 17 00:00:00 2001 From: Gabriel Majeri Date: Sat, 10 Aug 2019 08:10:56 +0300 Subject: [PATCH 1/4] Add Keras layer for using compressed embeddings --- nncompress/__init__.py | 1 + nncompress/keras.py | 80 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 nncompress/keras.py diff --git a/nncompress/__init__.py b/nncompress/__init__.py index dc6ea56..9de0b7c 100644 --- a/nncompress/__init__.py +++ b/nncompress/__init__.py @@ -1,3 +1,4 @@ from __future__ import absolute_import, division, print_function from .embed_compress import EmbeddingCompressor +from .keras import CompressedEmbedding diff --git a/nncompress/keras.py b/nncompress/keras.py new file mode 100644 index 0000000..e064ac0 --- /dev/null +++ b/nncompress/keras.py @@ -0,0 +1,80 @@ +import numpy as np +from keras import backend as K +from keras.layers import Layer +from keras.utils.generic_utils import to_list + +class CompressedEmbedding(Layer): + """Embedding layer which uses compressed embeddings.""" + + def __init__(self, codebook, codes, input_length=None, **kwargs): + """Initializes a new compressed embedding layer. + + - `codebook` is a matrix of codebooks which map indices to basis vectors + - `codes` is a matrix which maps word indices to sequences of integers, + representing indexes into each codebook. + """ + + if 'input_shape' not in kwargs: + if input_length: + kwargs['input_shape'] = (input_length,) + else: + kwargs['input_shape'] = (None,) + + assert isinstance(codebook, np.ndarray) + assert isinstance(codes, np.ndarray) + + self.codebook_np = codebook + self.codes_np = codes + + self.input_length = input_length + self.output_dim = codebook.shape[-1] + + super().__init__(**kwargs) + + def build(self, input_shape): + self.codebook = K.constant(self.codebook_np, dtype='float32', name='codebook') + self.codes = K.constant(self.codes_np, dtype='int32', name='word_codes') + + super().build(input_shape) + + def call(self, x): + if K.dtype(x) != 'int32': + x = K.cast(x, 'int32') + + # Get the indices into the codebooks + codes = K.gather(self.codes, x) + + # Gather the required basis vectors for these words + vectors = K.gather(self.codebook, codes) + + # Sum the basis vectors to obtain the embedding vectors + embeddings = K.sum(vectors, axis=-2) + + return embeddings + + def compute_output_shape(self, input_shape): + """Computes the output shape of this embedding layer. + + Code taken from the original Keras `Embedding` layer. + """ + + if self.input_length is None: + return input_shape + (self.output_dim,) + + # input_length can be tuple if input is 3D or higher + in_lens = to_list(self.input_length, allow_tuple=True) + if len(in_lens) != len(input_shape) - 1: + raise ValueError( + '"input_length" is %s, but received input has shape %s' % + (str(self.input_length), str(input_shape))) + + for i, (s1, s2) in enumerate(zip(in_lens, input_shape[1:])): + if s1 is not None and s2 is not None and s1 != s2: + raise ValueError( + '"input_length" is %s, but received input has shape %s' % + (str(self.input_length), str(input_shape))) + + if s1 is None: + in_lens[i] = s2 + + return (input_shape[0],) + tuple(in_lens) + (self.output_dim,) From 77209acbd1efd820b449586e47c4ea0614aacd1a Mon Sep 17 00:00:00 2001 From: Gabriel Majeri Date: Sat, 10 Aug 2019 08:18:39 +0300 Subject: [PATCH 2/4] Store codebook and code vectors in `__init__` --- nncompress/keras.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/nncompress/keras.py b/nncompress/keras.py index e064ac0..2756351 100644 --- a/nncompress/keras.py +++ b/nncompress/keras.py @@ -23,8 +23,8 @@ def __init__(self, codebook, codes, input_length=None, **kwargs): assert isinstance(codebook, np.ndarray) assert isinstance(codes, np.ndarray) - self.codebook_np = codebook - self.codes_np = codes + self.codebook = K.constant(codebook, dtype='float32', name='codebook') + self.codes = K.constant(codes, dtype='int32', name='word_codes') self.input_length = input_length self.output_dim = codebook.shape[-1] @@ -32,9 +32,6 @@ def __init__(self, codebook, codes, input_length=None, **kwargs): super().__init__(**kwargs) def build(self, input_shape): - self.codebook = K.constant(self.codebook_np, dtype='float32', name='codebook') - self.codes = K.constant(self.codes_np, dtype='int32', name='word_codes') - super().build(input_shape) def call(self, x): From 831ca0f2f9085e895ab77dc3c4f051d7020ecde1 Mon Sep 17 00:00:00 2001 From: Gabriel Majeri Date: Wed, 14 Aug 2019 10:52:53 +0300 Subject: [PATCH 3/4] Fix decompression, use TensorFlow directly --- nncompress/keras.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/nncompress/keras.py b/nncompress/keras.py index 2756351..a0605e8 100644 --- a/nncompress/keras.py +++ b/nncompress/keras.py @@ -1,4 +1,7 @@ import numpy as np + +import tensorflow as tf + from keras import backend as K from keras.layers import Layer from keras.utils.generic_utils import to_list @@ -22,6 +25,7 @@ def __init__(self, codebook, codes, input_length=None, **kwargs): assert isinstance(codebook, np.ndarray) assert isinstance(codes, np.ndarray) + assert len(codebook.shape) == 3 self.codebook = K.constant(codebook, dtype='float32', name='codebook') self.codes = K.constant(codes, dtype='int32', name='word_codes') @@ -35,17 +39,19 @@ def build(self, input_shape): super().build(input_shape) def call(self, x): - if K.dtype(x) != 'int32': - x = K.cast(x, 'int32') + x = tf.cast(x, tf.int32) # Get the indices into the codebooks - codes = K.gather(self.codes, x) + codes = tf.gather(self.codes, x) + + indices = tf.broadcast_to(tf.range(self.codebook.shape[0]), codes.shape) + indices = tf.stack([indices, codes], axis=-1) # Gather the required basis vectors for these words - vectors = K.gather(self.codebook, codes) + vectors = tf.gather_nd(self.codebook, indices) # Sum the basis vectors to obtain the embedding vectors - embeddings = K.sum(vectors, axis=-2) + embeddings = tf.reduce_sum(vectors, axis=-2) return embeddings From 8f1be04ce17f58e9b6e74479bff792a9a6549942 Mon Sep 17 00:00:00 2001 From: Gabriel Majeri Date: Wed, 14 Aug 2019 11:09:27 +0300 Subject: [PATCH 4/4] Fix TensorFlow error with non-eager mode --- nncompress/keras.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/nncompress/keras.py b/nncompress/keras.py index a0605e8..404dbe5 100644 --- a/nncompress/keras.py +++ b/nncompress/keras.py @@ -26,6 +26,8 @@ def __init__(self, codebook, codes, input_length=None, **kwargs): assert isinstance(codebook, np.ndarray) assert isinstance(codes, np.ndarray) assert len(codebook.shape) == 3 + + self.codebook_count = codebook.shape[0] self.codebook = K.constant(codebook, dtype='float32', name='codebook') self.codes = K.constant(codes, dtype='int32', name='word_codes') @@ -42,16 +44,16 @@ def call(self, x): x = tf.cast(x, tf.int32) # Get the indices into the codebooks - codes = tf.gather(self.codes, x) + codes = K.gather(self.codes, x) - indices = tf.broadcast_to(tf.range(self.codebook.shape[0]), codes.shape) + indices = tf.broadcast_to(tf.range(self.codebook_count), tf.shape(codes)) indices = tf.stack([indices, codes], axis=-1) # Gather the required basis vectors for these words vectors = tf.gather_nd(self.codebook, indices) # Sum the basis vectors to obtain the embedding vectors - embeddings = tf.reduce_sum(vectors, axis=-2) + embeddings = K.sum(vectors, axis=-2) return embeddings