Infer dtype in SymbolBlock import from input symbol

python/mxnet/gluon/block.py

@@ -1053,13 +1053,34 @@ def __init__(self, outputs, inputs, params=None):
                     "SymbolBlock doesn't support Parameter '%s' because its storage " \
                     "type is 'row_sparse'." % j.name
 
-        for i in out.list_arguments():
-            if i not in input_names:
-                self.params.get(i, allow_deferred_init=True)
+        # Infer type of parameters. Without this, every parameter will be created with
+        # default type i.e., fp32
+        arg_params = out.list_arguments()
+        aux_params = out.list_auxiliary_states()
+
+        infer_type_success, arg_types, aux_types = _infer_param_types(inputs[0],
+                                                                      out,
+                                                                      arg_params,
+                                                                      aux_params)
+
+        if infer_type_success:
+            # Use inferred types for params
+            for i, arg in enumerate(arg_params):
+                if arg not in input_names:
+                    self.params.get(arg, allow_deferred_init=True, dtype=arg_types[i])
+
+            for i, aux in enumerate(aux_params):
+                if aux not in input_names:
+                    self.params.get(aux, grad_req='null', allow_deferred_init=True, dtype=aux_types[i])
+        else:
+            # Use default types for params
+            for i, arg in enumerate(arg_params):
+                if arg not in input_names:
+                    self.params.get(arg, allow_deferred_init=True)
 
-        for i in out.list_auxiliary_states():
-            if i not in input_names:
-                self.params.get(i, grad_req='null', allow_deferred_init=True)
+            for i, aux in enumerate(aux_params):
+                if aux not in input_names:
+                    self.params.get(aux, grad_req='null', allow_deferred_init=True)
 
         self._cached_graph = syms, out
         len_prefix = len(_common_prefix(list(self._params.keys())))
@@ -1086,3 +1107,49 @@ def _clear_cached_op(self):
 
     def hybrid_forward(self, F, x, *args, **kwargs):
         raise NotImplementedError
+
+def _infer_param_types(in_params, out_params, arg_params, aux_params):
+    """Utility function that helps in inferring DType of args and auxs params
+    from given input param.
+
+    Parameters
+    ----------
+    in_params: Symbol
+        Input symbol variable.
+    out_params: Symbol
+        Output symbol variable.
+    arg_params: List of Str
+        List of names of argument parametrs.
+    aux_params: List of Str
+        List of names of auxiliary parameters.
+
+    Returns
+    -------
+    infer_type_success: Boolean
+        True if able to infer types for all given arg_params and aux_params.
+        False, otherwise.
+    arg_types: List of numpy.dtype
+        List of arg_params type. Order is same as arg_params.
+        None if unable to infer type.
+    aux_types: List of numpy.dtype
+        List of aux_params type. Order is same as aux_params.
+        None if unable to infer type.
+    """
+    infer_type_success = False
+    arg_types = None
+    aux_types = None
+
+    # Get Input symbol details. This will be used to infer types of
+    # other parameters.
+    input_sym_name = in_params.name
+    input_sym_arg_type = in_params.infer_type()[0]
+
+    # Try to infer types of other parameters.
+    if input_sym_arg_type and len(input_sym_arg_type) > 0:
+        params = {input_sym_name:input_sym_arg_type[0]}
+        arg_types, _, aux_types = out_params.infer_type(**params)
+        if arg_types is not None and len(arg_types) == len(arg_params) and \
+           aux_types is not None and len(aux_types) == len(aux_params):
+            infer_type_success = True
+
+    return (infer_type_success, arg_types, aux_types)

tests/python/unittest/test_gluon.py

@@ -15,6 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import os
+import tempfile
+
 import mxnet as mx
 from mxnet import gluon
 from mxnet.gluon import nn
@@ -336,6 +339,28 @@ def hybrid_forward(self, F, x):
     net.hybridize()
     assert isinstance(net(mx.nd.zeros((16, 10))), mx.nd.NDArray)
 
+    # Test case to verify if initializing the SymbolBlock from a model with params 
+    # other than fp32 param dtype.
+    # Load a resnet model, cast it to fp64 and export
+    tmp = tempfile.mkdtemp()
+    tmpfile = os.path.join(tmp, 'resnet34_fp64')
+    ctx = mx.cpu(0)
+
+    net_fp32 = mx.gluon.model_zoo.vision.resnet34_v2(pretrained=True, ctx=ctx)
+    net_fp32.cast('float64')
+    net_fp32.hybridize()
+    data = mx.nd.zeros((1,3,224,224), dtype='float64', ctx=ctx)
+    net_fp32.forward(data)
+    net_fp32.export(tmpfile, 0)
+
+    # Load the saved model and verify if all the params are loaded correctly.
+    # and choose one of the param to verify the type.
+    sm = mx.sym.load(tmpfile + '-symbol.json')
+    inputs = mx.sym.var('data', dtype='float64')
+    net_fp64 = mx.gluon.SymbolBlock(sm, inputs)
+    net_fp64.collect_params().load(tmpfile + '-0000.params', ctx=ctx)
+    assert (net_fp64.params['resnetv20_stage1_conv2_weight'].dtype is np.float64)
+
 @with_seed()
 @raises(AssertionError)
 def test_sparse_symbol_block():