Repository: sql-machine-learning/elasticdl

Issue:
Concatenate IDs of multiple category features for embedding.
### Why do we need to concatenate IDs of multiple category features for embedding?
In #1721, we have introduced that we need to convert categorical feature values to integer IDs if we want to export the model using Tensorflow SavedModel for TF Serving.

|  age | education | marital-status |
| ---- | --- | --- |
|  34  | Master | Divorced |
|  54  | Doctor | Never-married |
|  42  | Bachelor | Never-married |

To

|  age | education | marital-status |
| ---- | --- | --- |
|  34  | 0 | 0 |
|  54  | 1 | 1 |
|  42  | 2 | 1 |

After converting category values to IDs, we generally use those IDs to perform a lookup in the embedding matrix. 

**The problem: a dataset sometimes has many categorical features. If we make embedding for each categorical features separately, we need to create many embedding table variables. Besides the weights in variables, there are overhead to create a variable. So, the size of the model may be very huge and the performance of embedding lookup may be inefficient.**
![image](https://user-images.githubusercontent.com/18071380/74301254-795bac80-4d8d-11ea-98e3-2632e47df669.png)

In order to reduce the number of embedding table variables, we can concatenate the categorical feature IDs tensor to a big tensor and merge the embedding tables. However, the same IDs will return the same embedding vectors by lookup in the merged embedding table. In the following figure, we can see that embedding vectors of "marital-status" are the same as "education".
![image](https://user-images.githubusercontent.com/18071380/74301128-00f4eb80-4d8d-11ea-97eb-7db2e798ffb2.png)

So, we need to add an offset for IDs of "marital-status" so that the "martial-status" feature can get its embedding vectors by lookup in the merged embedding table. 
![image](https://user-images.githubusercontent.com/18071380/74301135-06eacc80-4d8d-11ea-9f2b-2a577ed52b0f.png)



### Solution Proposals to Concatenate the IDs using different Tensorflow API.
#1721 has listed 3 methods to convert category value to IDs. We need to adopt different methods to concatenate IDs using different methods.

#### 1. Concatenate the IDs generated by categorical columns in `tf.feature_column` such as `tf.feature_column.categorical_column_with_hash_bucket`. 

The example is the 1st case showed in #1721. If we use categorical columns to convert categories values to IDs, we must use `embedding_column` to make embedding for those IDs. Because the output of categorial columns is a sparse tensor which cannot be directly used in `DenseFeature`. So, we need to concatenate the outputs of categorial features before `embedding_column` to reduce the number of embedding variables. 

```python
education_hash_column = tf.feature_column.categorical_column_with_hash_bucket(
	name="education", hash_bucket_size=3
)  # the id is in [0,3)
marital_hash_column = tf.feature_column.categorical_column_with_hash_bucket(
	name="marital-status", hash_bucket_size=5
) # the id is in [0,5)

edu_marital_concat = concat_column([education_hash_column, marital_hash_column])

education_embedded_column = tf.feature_column.embedding_column(
	edu_marital_concat, embedding_dim=2
)
```

The `concat_column` will concatenate the outputs of `education_hash_column` and `marital_hash_column` and add offset 3 for the IDs of `marital_hash_column`.  

**In the case, we need to customize a `concat_column` showed in PR #1719.**

#### 2. Concatenate the IDs generated by `numeric_column` with custom transform_fn。

The examples are 2nd  in #1721. The output of `numeric_column` is a tensor with IDs and we can directly use it in `DenseFeatures`.  

```python
def generate_hash_bucket_column(name, hash_bucket_size):
    def hash_bucket_id(x, hash_bucket_size):
        if x.dtype is not tf.string:
            x = tf.strings.as_string(x)
        return tf.strings.to_hash_bucket_fast(x, hash_bucket_size)

    transform_fn = lambda x, hash_bucket_size=hash_bucket_size : (
        hash_bucket_id(x, hash_bucket_size)
    )
    return tf.feature_column.numeric_column(
        name, dtype=tf.int32, normalizer_fn=transform_fn
    )

input_layers = [
	tf.keras.layers.Input(name="education", shape=(1,), dtype=tf.string)
	tf.keras.layers.Input(name="marital-status", shape=(1,), dtype=tf.string)
]
education_hash = generate_hash_bucket_column(
	name="education", hash_bucket_size=3
) # the id is in [0,3)
marital_hash = generate_hash_bucket_column(
	name="marital-status", hash_bucket_size=5
) # the id is in [0,5)

education_hash_ids = tf.keras.layers.DenseFeature([education_hash])(input_layers)
marital_hash_ids = tf.keras.layers.DenseFeature([marital_hash])(input_layers)
```
Then,  we can add offset for the IDs tensor of "marital-status" and concatenate it with "education" IDs tensor like:
```python
marital_ids_with_offset = marital_hash_ids + 3 #3 is the number of education IDs
edu_marital_concat = tf.keras.layers.Concatenate()([education_hash_ids, marital_ids_with_offset])
```

**In those case, we need to customize a transform_fn for `numeric_column` and don't need to customize a `concat_column`.**

#### 3. Concatenate the IDs generated by custom transformation layers.
The examples are the 3rd methods in #1721. The output of the custom layer `HashBucket` is the same as the `numeric_column` in the 2nd method. So we can use the same way to add offset to "marital-status" IDs and concatenate.

```python
class HashBucket(tf.keras.layers.Layer):
    def __init__(self, hash_bucket_size):
        super(HashBucket, self).__init__()
        self.hash_bucket_size =hash_ bucket_size

    def call(self, inputs):
        if inputs.dtype is not tf.string:
            inputs = tf.strings.as_string(inputs)
        bucket_id = tf.strings.to_hash_bucket_fast(
        	inputs, self.hash_bucket_size
        )
        return tf.cast(bucket_id, tf.int64)

education_input = tf.keras.layers.Input(name="education", shape=(1,), dtype=tf.string)
marital_input = tf.keras.layers.Input(name="marital-status", shape=(1,), dtype=tf.string)

education_hash_ids = HashBucket(hash_bucket_size=3)(education_input) # the id is in [0,3)
marital_hash_ids = HashBucket(hash_bucket_size=5)(marital_input) # the id is in [0,5)

marital_ids_with_offset = marital_hash_ids + 3 #3 is the number of education IDs

edu_marital_concat = tf.keras.layers.Concatenate()([education_hash_ids, marital_ids_with_offset])
```




Discussion:

Resolving patch:
diff --git a/elasticdl_preprocessing/layers/concatenate_with_offset.py b/elasticdl_preprocessing/layers/concatenate_with_offset.py
new file mode 100644
--- /dev/null
+++ b/elasticdl_preprocessing/layers/concatenate_with_offset.py
@@ -0,0 +1,66 @@
+import tensorflow as tf
+
+
+class ConcatenateWithOffset(tf.keras.layers.Concatenate):
+    """Layer that add offset to each id tensor in the input list and
+    then concatenate these tensors.
+
+    It takes as input a list of tensors and returns a single tensor.
+    Firstly, it will add an offset in offsets for each tensor in inputs.
+    Then concatenate them to a single tensor. The tensor in inputs
+    must have the same type, `Tensor` or `RaggedTensor` or `SparseTensor` and
+    the same shape.
+
+    Example :
+    ```python
+        a1 = tf.constant([[1], [1], [1]])
+        a2 = tf.constant([[2], [2], [2]])
+        offsets = [0, 10]
+        layer = ConcatenateWithOffset(offsets=offsets, axis=1)
+        layer([a1, a2])
+        [[ 1 12]
+         [ 1 12]
+         [ 1 12]]
+    ```
+
+    Arguments:
+        offsets: numeric list to add
+        axis: Axis along which to concatenate.
+        **kwargs: standard layer keyword arguments.
+    """
+
+    def __init__(self, offsets, axis=-1):
+        super(ConcatenateWithOffset, self).__init__()
+        self.offsets = offsets
+        self.axis = axis
+
+    def call(self, inputs):
+        ids_with_offset = []
+        if len(self.offsets) != len(inputs):
+            raise ValueError(
+                "The offsets length is not equal to inputs length"
+                "the inputs are {}, offsets are {}".format(
+                    inputs, self.offsets
+                )
+            )
+        for i, tensor in enumerate(inputs):
+            if isinstance(tensor, tf.SparseTensor):
+                ids_with_offset.append(
+                    tf.SparseTensor(
+                        indices=tensor.indices,
+                        values=tensor.values + self.offsets[i],
+                        dense_shape=tensor.dense_shape,
+                    )
+                )
+            else:
+                ids_with_offset.append(tensor + self.offsets[i])
+
+        if isinstance(ids_with_offset[0], tf.SparseTensor):
+            result = tf.sparse.concat(
+                axis=self.axis, sp_inputs=ids_with_offset
+            )
+        else:
+            result = tf.keras.layers.concatenate(
+                ids_with_offset, axis=self.axis
+            )
+        return result


Test patch:
diff --git a/elasticdl_preprocessing/tests/concatenate_with_offset_test.py b/elasticdl_preprocessing/tests/concatenate_with_offset_test.py
new file mode 100644
--- /dev/null
+++ b/elasticdl_preprocessing/tests/concatenate_with_offset_test.py
@@ -0,0 +1,40 @@
+import unittest
+
+import numpy as np
+import tensorflow as tf
+
+from elasticdl_preprocessing.layers.concatenate_with_offset import (
+    ConcatenateWithOffset,
+)
+from elasticdl_preprocessing.tests.test_utils import (
+    ragged_tensor_equal,
+    sparse_tensor_equal,
+)
+
+
+class ConcatenateWithOffsetTest(unittest.TestCase):
+    def test_concatenate_with_offset(self):
+        tensor_1 = tf.constant([[1], [1], [1]])
+        tensor_2 = tf.constant([[2], [2], [2]])
+        offsets = [0, 10]
+        concat_layer = ConcatenateWithOffset(offsets=offsets, axis=1)
+
+        output = concat_layer([tensor_1, tensor_2])
+        expected_out = np.array([[1, 12], [1, 12], [1, 12]])
+        self.assertTrue(np.array_equal(output.numpy(), expected_out))
+
+        ragged_tensor_1 = tf.ragged.constant([[1], [], [1]], dtype=tf.int64)
+        ragged_tensor_2 = tf.ragged.constant([[2], [2], []], dtype=tf.int64)
+        output = concat_layer([ragged_tensor_1, ragged_tensor_2])
+        expected_out = tf.ragged.constant([[1, 12], [12], [1]], dtype=tf.int64)
+        self.assertTrue(ragged_tensor_equal(output, expected_out))
+
+        sparse_tensor_1 = ragged_tensor_1.to_sparse()
+        sparse_tensor_2 = ragged_tensor_2.to_sparse()
+        output = concat_layer([sparse_tensor_1, sparse_tensor_2])
+        expected_out = tf.SparseTensor(
+            indices=np.array([[0, 0], [0, 1], [1, 1], [2, 0]]),
+            values=np.array([1, 12, 12, 1]),
+            dense_shape=(3, 2),
+        )
+        self.assertTrue(sparse_tensor_equal(output, expected_out))
diff --git a/elasticdl_preprocessing/tests/test_utils.py b/elasticdl_preprocessing/tests/test_utils.py
--- a/elasticdl_preprocessing/tests/test_utils.py
+++ b/elasticdl_preprocessing/tests/test_utils.py
@@ -1,9 +1,13 @@
 import numpy as np
 import tensorflow as tf
-from tensorflow.python.ops.ragged import ragged_tensor
 
 
 def sparse_tensor_equal(sp_a, sp_b):
+    if not isinstance(sp_a, tf.SparseTensor) or not isinstance(
+        sp_b, tf.SparseTensor
+    ):
+        return False
+
     if not np.array_equal(sp_a.dense_shape.numpy(), sp_b.dense_shape.numpy()):
         return False
 
@@ -20,15 +24,19 @@ def sparse_tensor_equal(sp_a, sp_b):
 
 
 def ragged_tensor_equal(rt_a, rt_b):
-    print(rt_a, rt_b)
+    if not isinstance(rt_a, tf.RaggedTensor) or not isinstance(
+        rt_b, tf.RaggedTensor
+    ):
+        return False
+
     if rt_a.shape.as_list() != rt_b.shape.as_list():
         return False
 
     for i in range(rt_a.shape[0]):
         sub_rt_a = rt_a[i]
         sub_rt_b = rt_b[i]
-        if ragged_tensor.is_ragged(sub_rt_a) and ragged_tensor.is_ragged(
-            sub_rt_b
+        if isinstance(sub_rt_a, tf.RaggedTensor) and isinstance(
+            sub_rt_b, tf.RaggedTensor
         ):
             if not ragged_tensor_equal(sub_rt_a, sub_rt_b):
                 return False