用于 tensorflow keras 模型训练的 SparseTensor 生成器

Question

我有大量数据适合模型训练。并且在模型构建中，我们使用了 tf.keras.experimental.SequenceFeatures，它要求输入应该是 SpareTensor。 我厌倦了使用 generaotr，但不适用于 SparsTensor。 这是示例代码：

from models.model_attention import AttentionModel
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import LSTM, Dropout, Dense


inputs = {'f1': tf.keras.layers.Input(name='f1', sparse=True, shape=(40, 1), dtype='float32'),
          'f2': tf.keras.layers.Input(name='f2', sparse=True, shape=(40, 1), dtype='float32')}

features = [tf.feature_column.sequence_numeric_column('f1', dtype=tf.float32),
            tf.feature_column.sequence_numeric_column('f2', dtype=tf.float32)]

input_layer, _ = tf.keras.experimental.SequenceFeatures(features)(inputs)
lstm_out = LSTM(128, return_sequences=False)(input_layer)
lstm_out = Dropout(0.2)(lstm_out)
lstm_out = Dense(1, activation='tanh')(lstm_out)
model = tf.keras.models.Model(inputs, lstm_out)
model.compile(loss='mse', metrics='mae', optimizer='Adam')


def gen():
    batch = 4
    while True:
        x1 = tf.sparse.from_dense(np.random.random((batch, 40, 1)))
        x2 = tf.sparse.from_dense(np.random.random((batch, 40, 1)))
        x = {'f1': x1, 'f2': x2}
        y = np.random.random((batch, 1))
        yield x, y


x, y = gen().__next__()
# x, y yielded from generator works
model.fit(x, y, epochs=2, verbose=2)
g = gen()
# TypeError: Input must be a SparseTensor.
model.fit(g, steps_per_epoch=2, epochs=2, verbose=2, validation_data=g, validation_steps=2)

gen() 函数用于正常的 numpy 数组生成器，但不适用于 SparseTensor 输入。 错误信息：

Traceback (most recent call last):
  File "D:/PycharmProjects/infinity_stock/tmp2.py", line 36, in <module>
    model.fit(g, steps_per_epoch=2, epochs=2, verbose=2, validation_data=g, validation_steps=2)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\training.py", line 1183, in fit
    tmp_logs = self.train_function(iterator)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\def_function.py", line 889, in __call__
    result = self._call(*args, **kwds)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\def_function.py", line 917, in _call
    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\function.py", line 3022, in __call__
    filtered_flat_args) = self._maybe_define_function(args, kwargs)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\function.py", line 3440, in _maybe_define_function
    return self._define_function_with_shape_relaxation(
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\function.py", line 3362, in _define_function_with_shape_relaxation
    graph_function = self._create_graph_function(
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\function.py", line 3279, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\framework\func_graph.py", line 999, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\eager\def_function.py", line 672, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\framework\func_graph.py", line 986, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:

    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\training.py:855 train_function  *
        return step_function(self, iterator)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\training.py:845 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:1285 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:2833 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:3608 _call_for_each_replica
        return fn(*args, **kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\training.py:838 run_step  **
        outputs = model.train_step(data)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\training.py:795 train_step
        y_pred = self(x, training=True)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\base_layer.py:1030 __call__
        outputs = call_fn(inputs, *args, **kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\functional.py:420 call
        return self._run_internal_graph(
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\functional.py:556 _run_internal_graph
        outputs = node.layer(*args, **kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\engine\base_layer.py:1030 __call__
        outputs = call_fn(inputs, *args, **kwargs)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\keras\feature_column\sequence_feature_column.py:159 call  **
        dense_tensor, sequence_length = column.get_sequence_dense_tensor(
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\feature_column\sequence_feature_column.py:442 get_sequence_dense_tensor
        dense_tensor = sparse_ops.sparse_tensor_to_dense(
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\ops\sparse_ops.py:1714 sparse_tensor_to_dense
        sp_input = _convert_to_sparse_tensor(sp_input)
    D:\anaconda3\envs\stock\lib\site-packages\tensorflow\python\ops\sparse_ops.py:72 _convert_to_sparse_tensor
        raise TypeError("Input must be a SparseTensor.")

    TypeError: Input must be a SparseTensor.

对这个问题有什么建议吗？

Answer 1

通过调试，终于查明原因： 在 dataset_ops.py 中，如果未指定输出签名，函数 from_generator() 默认将使用 tensorSpec：

  def from_generator(generator,
                     output_types=None,
                     output_shapes=None,
                     args=None,
                     output_signature=None):
    """Creates a `Dataset` whose elements are generated by `generator`.

    The `generator` argument must be a callable object that returns
    an object that supports the `iter()` protocol (e.g. a generator function).

    The elements generated by `generator` must be compatible with either the
    given `output_signature` argument or with the given `output_types` and
    (optionally) `output_shapes` arguments, whichever was specified.

    The recommended way to call `from_generator` is to use the
    `output_signature` argument. In this case the output will be assumed to
    consist of objects with the classes, shapes and types defined by
    `tf.TypeSpec` objects from `output_signature` argument:

    >>> def gen():
    ...   ragged_tensor = tf.ragged.constant([[1, 2], [3]])
    ...   yield 42, ragged_tensor
    >>>
    >>> dataset = tf.data.Dataset.from_generator(
    ...      gen,
    ...      output_signature=(
    ...          tf.TensorSpec(shape=(), dtype=tf.int32),
    ...          tf.RaggedTensorSpec(shape=(2, None), dtype=tf.int32)))
    >>>
    >>> list(dataset.take(1))
    [(<tf.Tensor: shape=(), dtype=int32, numpy=42>,
    <tf.RaggedTensor [[1, 2], [3]]>)]

    There is also a deprecated way to call `from_generator` by either with
    `output_types` argument alone or together with `output_shapes` argument.
    In this case the output of the function will be assumed to consist of
    `tf.Tensor` objects with the types defined by `output_types` and with the
    shapes which are either unknown or defined by `output_shapes`.

    Note: The current implementation of `Dataset.from_generator()` uses
    `tf.numpy_function` and inherits the same constraints. In particular, it
    requires the dataset and iterator related operations to be placed
    on a device in the same process as the Python program that called
    `Dataset.from_generator()`. The body of `generator` will not be
    serialized in a `GraphDef`, and you should not use this method if you
    need to serialize your model and restore it in a different environment.

    Note: If `generator` depends on mutable global variables or other external
    state, be aware that the runtime may invoke `generator` multiple times
    (in order to support repeating the `Dataset`) and at any time
    between the call to `Dataset.from_generator()` and the production of the
    first element from the generator. Mutating global variables or external
    state can cause undefined behavior, and we recommend that you explicitly
    cache any external state in `generator` before calling
    `Dataset.from_generator()`.

    Args:
      generator: A callable object that returns an object that supports the
        `iter()` protocol. If `args` is not specified, `generator` must take no
        arguments; otherwise it must take as many arguments as there are values
        in `args`.
      output_types: (Optional.) A (nested) structure of `tf.DType` objects
        corresponding to each component of an element yielded by `generator`.
      output_shapes: (Optional.) A (nested) structure of `tf.TensorShape`
        objects corresponding to each component of an element yielded by
        `generator`.
      args: (Optional.) A tuple of `tf.Tensor` objects that will be evaluated
        and passed to `generator` as NumPy-array arguments.
      output_signature: (Optional.) A (nested) structure of `tf.TypeSpec`
        objects corresponding to each component of an element yielded by
        `generator`.

    Returns:
      Dataset: A `Dataset`.
    """
    if not callable(generator):
      raise TypeError("`generator` must be callable.")

    if output_signature is not None:
      if output_types is not None:
        raise TypeError("`output_types` can not be used together with "
                        "`output_signature`")
      if output_shapes is not None:
        raise TypeError("`output_shapes` can not be used together with "
                        "`output_signature`")
      if not all(
          isinstance(_, type_spec.TypeSpec)
          for _ in nest.flatten(output_signature)):
        raise TypeError("All the elements of `output_signature` must be "
                        "`tf.TypeSpec` objects.")
    else:
      if output_types is None:
        raise TypeError("Either `output_signature` or `output_types` must "
                        "be specified")

    if output_signature is None:
      if output_shapes is None:
        output_shapes = nest.map_structure(
            lambda _: tensor_shape.TensorShape(None), output_types)
      else:
        output_shapes = nest.map_structure_up_to(output_types,
                                                 tensor_shape.as_shape,
                                                 output_shapes)
      output_signature = nest.map_structure_up_to(output_types,
                                                  tensor_spec.TensorSpec,
                                                  output_shapes, output_types)
    if all(
        isinstance(x, tensor_spec.TensorSpec)
        for x in nest.flatten(output_signature)):
      output_types = nest.pack_sequence_as(
          output_signature, [x.dtype for x in nest.flatten(output_signature)])
      output_shapes = nest.pack_sequence_as(
          output_signature, [x.shape for x in nest.flatten(output_signature)])

    if args is None:
      args = ()
    else:
      args = tuple(ops.convert_n_to_tensor(args, name="args"))

    generator_state = DatasetV2._GeneratorState(generator)

目前解决方案是使用 tf.data.Dataset.from_generator()，并明确指定输出签名。 示例如下：

data = tf.data.Dataset.from_generator(gen, output_signature=({'f1': SparseTensorSpec(TensorShape([4, 40, 1]), tf.float64),
                                                              'f2': SparseTensorSpec(TensorShape([4, 40, 1]), tf.float64)},
                                                             TensorSpec(shape=(4, 1), dtype=tf.float32)))
data.prefetch(1)
model.fit(data, steps_per_epoch=2, epochs=2, verbose=2)