Repository: OpenNMT/OpenNMT-tf

Issue:
Latest (1.14) OpenNMT-tf gives error on inferencing
Hi,
I have trained a model ParallelEncoder having two input files (words and their features). This far all has worked just fine, but now I get following error when trying to run inference:
Traceback (most recent call last):
  File "opennmt/bin/main.py", line 192, in <module>
    main()
  File "opennmt/bin/main.py", line 175, in main
    log_time=args.log_prediction_time)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/runner.py", line 329, in infer
    hooks=infer_hooks):
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 548, in predict
    input_fn, model_fn_lib.ModeKeys.PREDICT)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1023, in _get_features_from_input_fn
    result = self._call_input_fn(input_fn, mode)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1135, in _call_input_fn
    return input_fn(**kwargs)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/model.py", line 492, in <lambda>
    maximum_labels_length=maximum_labels_length)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/model.py", line 415, in _input_fn_impl
    length_fn=self._get_features_length)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/data.py", line 352, in inference_pipeline
    _key_func, _reduce_func, window_size=batch_size))
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1140, in apply
    dataset = transformation_func(self)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 117, in _apply_fn
    window_size_func)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 422, in __init__
    self._make_key_func(key_func, input_dataset)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/grouping.py", line 451, in _make_key_func
    "`key_func` must return a single tf.int64 scalar tensor.")
ValueError: `key_func` must return a single tf.int64 scalar tensor.

When I run infer on 1.10 it will work just fine on the model trained on 1.14. Error seems to be present also on 1.13.1



Discussion:
Hi,

Thanks for reporting. Do you confirm you are using the `--auto_config` flag?

To workaround the issue, try adding this in your configuration:

```yml
infer:
  bucket_width: 0
```
Great. This now works on "old" models. bucket_width 0 or 1 work just fine. When I'm trying to infer the new multi-source transformer mode I started getting another error:
2018-11-28 10:58:38.020515: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
INFO:tensorflow:Running local_init_op.
2018-11-28 10:58:38.257302: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
INFO:tensorflow:Done running local_init_op.
2018-11-28 10:58:38.349261: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
2018-11-28 10:58:38.580481: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
Traceback (most recent call last):
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1292, in _do_call
    return fn(*args)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1277, in _run_fn
    options, feed_dict, fetch_list, target_list, run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1365, in _call_tf_sessionrun
    run_metadata)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
	 [[{{node transformer/decoder/while/Reshape_70}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
	 [[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "opennmt/bin/main.py", line 192, in <module>
    main()
  File "opennmt/bin/main.py", line 175, in main
    log_time=args.log_prediction_time)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/runner.py", line 329, in infer
    hooks=infer_hooks):
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 568, in predict
    preds_evaluated = mon_sess.run(predictions)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 671, in run
    run_metadata=run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1148, in run
    run_metadata=run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1239, in run
    raise six.reraise(*original_exc_info)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/six.py", line 693, in reraise
    raise value
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1224, in run
    return self._sess.run(*args, **kwargs)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1296, in run
    run_metadata=run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1076, in run
    return self._sess.run(*args, **kwargs)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 887, in run
    run_metadata_ptr)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1110, in _run
    feed_dict_tensor, options, run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1286, in _do_run
    run_metadata)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1306, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
	 [[node transformer/decoder/while/Reshape_70 (defined at /home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py:88)  = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
	 [[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Caused by op 'transformer/decoder/while/Reshape_70', defined at:
  File "opennmt/bin/main.py", line 192, in <module>
    main()
  File "opennmt/bin/main.py", line 175, in main
    log_time=args.log_prediction_time)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/runner.py", line 329, in infer
    hooks=infer_hooks):
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 550, in predict
    features, None, model_fn_lib.ModeKeys.PREDICT, self.config)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1168, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/model.py", line 152, in _model_fn
    _, predictions = self._build(features, labels, params, mode, config=config)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/sequence_to_sequence.py", line 277, in _build
    return_alignment_history=True))
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/decoders/decoder.py", line 378, in dynamic_decode_and_search
    min_decode_length=minimum_length)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 613, in beam_search
    back_prop=False)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3274, in while_loop
    return_same_structure)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2994, in BuildLoop
    pred, body, original_loop_vars, loop_vars, shape_invariants)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2929, in _BuildLoop
    body_result = body(*packed_vars_for_body)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 539, in inner_loop
    i, alive_seq, alive_log_probs, states)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 428, in grow_topk
    lambda t: _unmerge_beam_dim(t, batch_size, beam_size), flat_states)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 347, in map_structure
    structure[0], [func(*x) for x in entries])
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 347, in <listcomp>
    structure[0], [func(*x) for x in entries])
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 428, in <lambda>
    lambda t: _unmerge_beam_dim(t, batch_size, beam_size), flat_states)
  File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 88, in _unmerge_beam_dim
    return tf.reshape(tensor, new_shape)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/gen_array_ops.py", line 6296, in reshape
    "Reshape", tensor=tensor, shape=shape, name=name)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
    op_def=op_def)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 488, in new_func
    return func(*args, **kwargs)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3274, in create_op
    op_def=op_def)
  File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1770, in __init__
    self._traceback = tf_stack.extract_stack()

InvalidArgumentError (see above for traceback): Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
	 [[node transformer/decoder/while/Reshape_70 (defined at /home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py:88)  = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
	 [[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

The model trains just fine (and it is a copy from example)

Resolving patch:
diff --git a/opennmt/utils/data.py b/opennmt/utils/data.py
--- a/opennmt/utils/data.py
+++ b/opennmt/utils/data.py
@@ -330,11 +330,11 @@ def _inject_index(index, x):
     return x
 
   def _key_func(x):
-    length = tf.cast(length_fn(x), tf.int32)
-    bucket_id = tf.constant(0, dtype=tf.int32)
+    length = length_fn(x)
+    bucket_id = tf.constant(0, dtype=tf.int64)
     if not isinstance(length, list):
-      bucket_id = tf.maximum(bucket_id, length // bucket_width)
-    return tf.cast(bucket_id, tf.int64)
+      bucket_id = tf.maximum(bucket_id, tf.cast(length, bucket_id.dtype) // bucket_width)
+    return bucket_id
 
   def _reduce_func(unused_key, dataset):
     return dataset.apply(batch_dataset(batch_size))


Test patch: