Fix for symbol block serialization issue (#582)

src/gluonts/support/util.py

@@ -33,12 +33,12 @@
 # Third-party imports
 import mxnet as mx
 import numpy as np
+from mxnet.gluon.block import _flatten
 
 # First-party imports
 from gluonts.core.serde import dump_json, load_json
 from gluonts.model.common import Tensor
 
-
 MXNET_HAS_ERF = hasattr(mx.nd, "erf")
 MXNET_HAS_ERFINV = hasattr(mx.nd, "erfinv")
 
@@ -176,6 +176,7 @@ def get_hybrid_forward_input_names(hb: mx.gluon.HybridBlock):
  return param_names[1:] # skip: F
 
 
+# noinspection PyProtectedMember
 def hybrid_block_to_symbol_block(
  hb: mx.gluon.HybridBlock, data_batch: List[mx.nd.NDArray]
 ) -> mx.gluon.SymbolBlock:
@@ -203,7 +204,11 @@ def hybrid_block_to_symbol_block(
  with tempfile.TemporaryDirectory(
  prefix="gluonts-estimator-temp-"
  ) as model_dir:
- num_inputs = len(data_batch)
+ # when importing, SymbolBlock has to know about the total number
+ # of input symbols, including nested Tensors
+ flat_data_batch, _ = _flatten(data_batch, "input")
+ num_inputs = len(flat_data_batch)
+
  model_dir_path = Path(model_dir)
  model_name = "gluonts-model"
 
@@ -220,6 +225,7 @@ def hybrid_block_to_symbol_block(
  return sb
 
 
+# noinspection PyProtectedMember
 def export_symb_block(
  hb: mx.gluon.HybridBlock, model_dir: Path, model_name: str, epoch: int = 0
 ) -> None:
@@ -239,6 +245,11 @@ def export_symb_block(
  model parameters.
  """
  hb.export(path=str(model_dir / model_name), epoch=epoch)
+ with (model_dir / f"{model_name}-in_out_format.json").open("w") as fp:
+ in_out_format = dict(
+ in_format=hb._in_format, out_format=hb._out_format
+ )
+ print(dump_json(in_out_format), file=fp)
 
 
 def import_symb_block(
@@ -269,14 +280,27 @@ def import_symb_block(
  else:
  input_names = [f"data{i}" for i in range(num_inputs)]
 
+ # FIXME: prevents mxnet from failing with empty saved parameters list
+ param_file: Optional[str] = str(
+ model_dir / f"{model_name}-{epoch:04}.params"
+ )
+ if not mx.nd.load(param_file):
+ param_file = None
+
  # FIXME: mx.gluon.SymbolBlock cannot infer float_type and uses default np.float32
  # FIXME: https://github.com/apache/incubator-mxnet/issues/11849
- return mx.gluon.SymbolBlock.imports(
+ sb = mx.gluon.SymbolBlock.imports(
  symbol_file=str(model_dir / f"{model_name}-symbol.json"),
  input_names=input_names,
- param_file=str(model_dir / f"{model_name}-{epoch:04}.params"),
+ param_file=param_file,
  ctx=mx.current_context(),
  )
+ with (model_dir / f"{model_name}-in_out_format.json").open("r") as fp:
+ formats = load_json(fp.read())
+ sb._in_format = formats["in_format"]
+ sb._out_format = formats["out_format"]
+
+ return sb
 
 
 def export_repr_block(

test/support/test_util.py

@@ -10,6 +10,8 @@
 # on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
+import itertools
+from typing import List
 
 # Third-party imports
 import mxnet as mx
@@ -18,6 +20,7 @@
 
 # First-party imports
 from gluonts.support import util
+from gluonts.model.common import Tensor
 
 
 @pytest.mark.parametrize("vec", [[[1, 2, 3, 4, 5], [10, 20, 30, 40, 50]]])
@@ -118,3 +121,73 @@ def test_erfinv() -> None:
  # Text np
  y_np = util.erfinv(np, x)
  assert np.allclose(y_np, y_scipy, rtol=1e-3)
+
+
+def sym_block_import_export_test_cases():
+ # single nested input
+ class TestBlock1(mx.gluon.HybridBlock):
+ def hybrid_forward(self, F, x1: Tensor, x2: List[Tensor]):
+ return F.broadcast_mul(x1, x2[0])
+
+ # multiple nested inputs
+ class TestBlock2(mx.gluon.HybridBlock):
+ def hybrid_forward(self, F, x1: Tensor, x2: List[Tensor]):
+ return F.broadcast_add(
+ F.broadcast_mul(x1, x2[0]), F.broadcast_mul(x1, x2[1])
+ )
+
+ # multiple nested inputs, and parameterized
+ class TestBlock3(mx.gluon.HybridBlock):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ with self.name_scope():
+ self.my_param = self.params.get(
+ "my_param",
+ shape=(1,),
+ init=mx.init.Constant(5),
+ allow_deferred_init=True,
+ )
+
+ def hybrid_forward(self, F, x1: Tensor, x2: List[Tensor], my_param):
+ y = F.broadcast_mul(x2[1], my_param)
+ return F.broadcast_add(F.broadcast_mul(x1, x2[0]), y)
+
+ # multiple nested inputs, parameterized, with sub-block
+ class TestBlock4(mx.gluon.HybridBlock):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ with self.name_scope():
+ self.my_param = self.params.get(
+ "my_param",
+ shape=(1,),
+ init=mx.init.Constant(5),
+ allow_deferred_init=True,
+ )
+ self.dense_layer = mx.gluon.nn.Dense(3)
+
+ def hybrid_forward(self, F, x1: Tensor, x2: List[Tensor], my_param):
+ y = self.dense_layer(F.broadcast_mul(x2[1], my_param))
+ return F.broadcast_add(F.broadcast_mul(x1, x2[0]), y)
+
+ return [TestBlock1, TestBlock2, TestBlock3, TestBlock4]
+
+
+@pytest.mark.parametrize(
+ ["block_type", "hybridize"],
+ itertools.product(sym_block_import_export_test_cases(), [True, False]),
+)
+def test_symb_block_export_import_nested_array(block_type, hybridize) -> None:
+ x1 = mx.nd.array([1, 2, 3])
+ x2 = [mx.nd.array([1, 5, 5]), mx.nd.array([2, 3, 3])]
+
+ my_block = block_type()
+ my_block.collect_params().initialize()
+ if hybridize:
+ my_block.hybridize()
+ my_block(x1, x2)
+
+ sb = util.hybrid_block_to_symbol_block(my_block, [x1, x2])
+
+ assert np.allclose(sb(x1, x2).asnumpy(), my_block(x1, x2).asnumpy())