[incubator-mxnet] 05/10: [WIP] [MXNET-334] fix dtype of Embedding and add dtype argument to Dense (#10615)

[jx...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jxie pushed a commit to branch v1.2.0
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git

commit 450067e19824e6ec873c6013e7f630d4adc41138
Author: Ziyue Huang <zy...@gmail.com>
AuthorDate: Tue Apr 24 03:23:41 2018 +0800

    [WIP] [MXNET-334] fix dtype of Embedding and add dtype argument to Dense (#10615)
    
    * add dtype to Dense and Embedding
    
    * fix test_parameter_str
---
 python/mxnet/gluon/nn/basic_layers.py | 10 ++++++----
 tests/python/unittest/test_gluon.py   | 20 +++++++++++++++++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/python/mxnet/gluon/nn/basic_layers.py b/python/mxnet/gluon/nn/basic_layers.py
index 7641509..9201dc5 100644
--- a/python/mxnet/gluon/nn/basic_layers.py
+++ b/python/mxnet/gluon/nn/basic_layers.py
@@ -153,6 +153,8 @@ class Dense(HybridBlock):
         If true, all but the first axis of input data are collapsed together.
         If false, all but the last axis of input data are kept the same, and the transformation
         applies on the last axis.
+    dtype : str or np.dtype, default 'float32'
+        Data type of output embeddings.
     weight_initializer : str or `Initializer`
         Initializer for the `kernel` weights matrix.
     bias_initializer: str or `Initializer`
@@ -179,7 +181,7 @@ class Dense(HybridBlock):
           `(x1, x2, ..., xn, units)`.
     """
     def __init__(self, units, activation=None, use_bias=True, flatten=True,
-                 weight_initializer=None, bias_initializer='zeros',
+                 dtype='float32', weight_initializer=None, bias_initializer='zeros',
                  in_units=0, **kwargs):
         super(Dense, self).__init__(**kwargs)
         self._flatten = flatten
@@ -187,11 +189,11 @@ class Dense(HybridBlock):
             self._units = units
             self._in_units = in_units
             self.weight = self.params.get('weight', shape=(units, in_units),
-                                          init=weight_initializer,
+                                          init=weight_initializer, dtype=dtype,
                                           allow_deferred_init=True)
             if use_bias:
                 self.bias = self.params.get('bias', shape=(units,),
-                                            init=bias_initializer,
+                                            init=bias_initializer, dtype=dtype,
                                             allow_deferred_init=True)
             else:
                 self.bias = None
@@ -379,7 +381,7 @@ class Embedding(HybridBlock):
         self._kwargs = {'input_dim': input_dim, 'output_dim': output_dim,
                         'dtype': dtype}
         self.weight = self.params.get('weight', shape=(input_dim, output_dim),
-                                      init=weight_initializer,
+                                      init=weight_initializer, dtype=dtype,
                                       allow_deferred_init=True)
 
     def hybrid_forward(self, F, x, weight):
diff --git a/tests/python/unittest/test_gluon.py b/tests/python/unittest/test_gluon.py
index 57e5087..0a5bda8 100644
--- a/tests/python/unittest/test_gluon.py
+++ b/tests/python/unittest/test_gluon.py
@@ -126,7 +126,7 @@ def test_parameter_str():
     assert lines[0] == 'net1_ ('
     assert 'net1_dense0_weight' in lines[1]
     assert '(10, 5)' in lines[1]
-    assert 'numpy.float32' in lines[1]
+    assert 'float32' in lines[1]
     assert lines[2] == ')'
 
 
@@ -776,6 +776,24 @@ def test_dtype():
 
     mx.nd.waitall()
 
+    class Net(gluon.Block):
+        def __init__(self, in_dim, output_dim):
+            super(Net, self).__init__()
+            with self.name_scope():
+                self.embed = gluon.nn.Embedding(input_dim=in_dim, output_dim=output_dim,dtype=np.float64)
+                self.dense = gluon.nn.Dense(2, dtype=np.float64)
+
+        def forward(self, x):
+            e = self.embed(x)
+            assert(e.dtype == np.float64)
+            y = self.dense(e)
+            assert(y.dtype == np.float64)
+            return y
+
+    net = Net(5, 10)
+    net.initialize()
+    out = net(mx.nd.ones((3,), dtype=np.float64))
+    mx.nd.waitall()
 
 @with_seed()
 def test_fill_shape_load():

-- 
To stop receiving notification emails like this one, please contact
jxie@apache.org.