[KERAS]Embedding layer (apache#5444)

trevor-m · Jun 8, 2020 · 429d69f · 429d69f
1 parent 70f7f5b
commit 429d69f
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/python/tvm/relay/frontend/keras.py b/python/tvm/relay/frontend/keras.py
@@ -207,6 +207,14 @@ def _convert_permute(inexpr, keras_layer, _):
     return _op.transpose(inexpr, axes=(0,) + keras_layer.dims)
 
 
+def _convert_embedding(inexpr, keras_layer, etab):
+    indices = inexpr
+    weightList = keras_layer.get_weights()
+    weight = etab.new_const(weightList[0])
+    out = _op.take(weight, indices.astype('int32'), axis=0)
+
+    return out
+
 def _convert_dense(inexpr, keras_layer, etab):
     weightList = keras_layer.get_weights()
     weight = etab.new_const(weightList[0].transpose([1, 0]))
@@ -893,7 +901,7 @@ def _default_skip(inexpr, keras_layer, _): # pylint: disable=unused-argument
     'Maximum'                  : _convert_merge,
     'Dot'                      : _convert_merge,
     'Permute'                  : _convert_permute,
-    # 'Embedding'              : _convert_embedding,
+    'Embedding'                : _convert_embedding,
     # 'RepeatVector'           : _convert_repeat_vector,
 
     'InputLayer'               : _default_skip,

diff --git a/tests/python/frontend/keras/test_forward.py b/tests/python/frontend/keras/test_forward.py
@@ -466,6 +466,24 @@ def test_forward_zero_padding3d(self, keras):
             keras_model = keras.models.Model(data, x)
             verify_keras_frontend(keras_model, layout='NDHWC')
 
+
+    def test_forward_embedding(self, keras):
+        data = keras.layers.Input(shape=(2, 4), dtype="int32")
+        x = keras.layers.Embedding(10, 3)(data)
+        keras_model = keras.models.Model(data, x)
+        verify_keras_frontend(keras_model, need_transpose=False)
+
+        data = keras.layers.Input(shape=(2, 3, 4), dtype="int32")
+        x = keras.layers.Embedding(4, 5)(data)
+        keras_model = keras.models.Model(data, x)
+        verify_keras_frontend(keras_model, need_transpose=False)
+
+        data = keras.layers.Input(shape=(6, 2, 3, 4), dtype="int32")
+        x = keras.layers.Embedding(4, 5)(data)
+        keras_model = keras.models.Model(data, x)
+        verify_keras_frontend(keras_model, need_transpose=False)
+
+
 if __name__ == '__main__':
     for k in [keras, tf_keras]:
         sut = TestKeras()
@@ -497,4 +515,4 @@ def test_forward_zero_padding3d(self, keras):
         sut.test_forward_pool3d(keras=k)
         sut.test_forward_upsample3d(keras=k)
         sut.test_forward_zero_padding3d(keras=k)
-
+        sut.test_forward_embedding(keras=k)