Prediction from tf.data.Dataset

TensorFlow.js version

latest version

Node version

$ node --version
v12.13.1

Describe the problem or feature request

In my python code I have a tf.data.Dataset where a file list is mapped to a tf.py_function:

 dataset = tf.data.Dataset.from_tensor_slices({
        'audio_id': list(filenames),
        'start': list(starts),
        'end': list(ends)
    })
    dataset = dataset.map(
        lambda sample: dict(
            sample,
            **audio_adapter.load_tf_waveform(
                sample['audio_id'],
                session=session,
                sample_rate=sample_rate,
                offset=sample['start'],
                duration=sample['end'] - sample['start'])),
        num_parallel_calls=2)

When I load the model into tfjs with the new tf.node.loadSavedModel I get the model with the following signature:

TFSavedModel {
  sessionId: 0,
  jsid: 0,
  inputNodeNames: {
    audio_id: 'Placeholder_1:0',
    mix_spectrogram: 'strided_slice_3:0',
    mix_stft: 'transpose_1:0',
    waveform: 'Placeholder:0'
  },
  outputNodeNames: {
    accompaniment: 'strided_slice_23:0',
    audio_id: 'Placeholder_1:0',
    vocals: 'strided_slice_13:0'
  },
  backend: NodeJSKernelBackend {
    binding: {},
    isGPUPackage: false,
    isUsingGpuDevice: false,
    tensorMap: DataStorage {
      backend: [Circular],
      dataMover: [Engine],
      data: [WeakMap],
      dataIdsCount: 0
    }
  },
  disposed: false
}

When in python the predict takes the dataset as input:

fn = lambda: get_dataset(
                    audio_adapter,
                    filenames_and_crops,
                    sample_rate,
                    n_channels, session)
with session.as_default():
    with session.graph.as_default():
        prediction = estimator.predict(fn,yield_single_examples=False)

How load a this dataset format into tfjs model.predict?

Code to reproduce the bug / link to feature request

[danwexler]

Is there a feature request to add tf.data.Dataset as a supported input type for model.predict? In the meantime, I guess we use tf.stack and handle pre-loading ourselves? Is it worth looking at the Keras model.predict as an example of how to implement tf.data.Dataset support in model.predict?

[tafsiri]

Currently model.predict does not take a dataset but only takes tensors. You could however use functions like map custom loop to call the predict function of your loaded model with a realized tensor. You can use the other dataset methods to set up a pipeline to stream the data from disk. The generator function is probably the most flexible way to create a dataset from file streams (we don’t have a built in utility to stream data from files in node).

cc @kangyizhang in case i missed anything/got anything wrong.