l2normalization operator support for tensorflow

Author: siju-samuel

nnvm/python/nnvm/frontend/tensorflow.py

@@ -468,6 +468,21 @@ def _impl(inputs, attr, params):
             ignores=['index_type', 'T'])(new_inputs, attr)
     return _impl
 
+def _sum():
+    def _impl(inputs, attr, params):
+        axis = params.pop(inputs[1].list_output_names()[0]).asnumpy()
+        return AttrCvt(
+            op_name='sum',
+            extras={'axis': axis},
+            transforms={'keep_dims':'keepdims'},
+            ignores=['name', 'Tidx'])(inputs[0], attr)
+    return _impl
+
+def _square():
+    def _impl(inputs, attr, params):
+        return _sym.elemwise_mul(inputs[0], inputs[0])
+    return _impl
+
 def _gather_v2():
     "Tensorflow now support only gatherv2"
     def _impl(inputs, attr, params):
@@ -663,13 +678,17 @@ def _impl(inputs, in_state_c, in_state_h, attr, params):
     'Identity'                          : _identity(),
     'MatMul'                            : _matmul(),
     'MaxPool'                           : _pooling('max_pool'),
+    'Add'                               : _elemwise('add'),
+    'Sub'                               : _elemwise('sub'),
     'Mul'                               : _elemwise('mul'),
+    'Maximum'                           : _elemwise('max'),
+    'Minimum'                           : _elemwise('min'),
+    'Sum'                               : _sum(),
+    'Square'                            : _square(),
     'Relu'                              : AttrCvt('relu'),
     'Reshape'                           : _reshape(),
     'ResizeBilinear'                    : _resize_bilinear(),
     'Softmax'                           : AttrCvt('softmax', {'axis': ('axis', 1)}),
-    'Sub'                               : _elemwise('sub'),
-    'Add'                               : _elemwise('add'),
     'Rsqrt'                             : _rsqrt(),
     'Squeeze'                           : _squeeze(),
     'FusedBatchNorm'                    : _fused_batch_norm(),

nnvm/tests/python/frontend/tensorflow/test_forward.py

@@ -12,6 +12,7 @@
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import graph_util
 from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import nn
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import math_ops
@@ -854,6 +855,41 @@ def _get_sample(data, state):
         np.testing.assert_allclose(tf_samples, tvm_samples, rtol=1e-5, atol=1e-5)
         assert(tvm_sample_str == tf_sample_str)
 
+#######################################################################
+# l2_normalize
+# ------------
+def _test_l2_normalize(ishape, eps, axis):
+    """ testing l2 normalize (uses max, sum, square, sqrt frontend operators)"""
+
+    inp_array = np.random.uniform(size=ishape).astype(np.float32)
+    inp_array.fill(1)
+
+    with tf.Graph().as_default():
+        in1 = tf.placeholder(shape=inp_array.shape, dtype=inp_array.dtype, name="Placeholder")
+        nn.l2_normalize(in1,
+                        axis=axis,
+                        epsilon=eps,
+                        name=None,
+                        dim=None)
+
+        with tf.Session() as sess:
+            graph_def = tf.graph_util.convert_variables_to_constants(
+                sess,
+                sess.graph.as_graph_def(add_shapes=True),
+                ['l2_normalize'],
+                )
+            tf_output = run_tf_graph(sess, inp_array, 'Placeholder:0', 'Placeholder:0')
+            tvm_output = run_tvm_graph(graph_def,
+                                       inp_array,
+                                       "Placeholder",
+                                       tf_output.shape,
+                                       tf_output.dtype)
+
+            np.testing.assert_allclose(tf_output, tvm_output, atol=1e-3, rtol=1e-3)
+            sess.close()
+def test_forward_l2_normalize():
+    _test_l2_normalize((1, 3, 20, 20), 0.001, (0,))
+
 #######################################################################
 # Main
 # ----
@@ -875,3 +911,4 @@ def _get_sample(data, state):
     test_forward_stridedslice()
     test_forward_gather()
     test_forward_ptb()
+    test_forward_l2_normalize()