Tensorflow：如何在 tf.data.Dataset from_generator 中保留文件名？

Question

I am struggling with the following. I am creating a tf.data.Dataset using the from_generator method. It works great, but after prediction I would like to investigate which samples were misclassified and why. For that, I need to retrieve the file names in the same order they were fed to a model. How can I do it?

def make_dataset(directory):
    """ Makes a dataset from generator. """

    def generator():
        files = list_files_in_directory(directory)
        random.Random(42).shuffle(files)
        print(f'Files in {directory}: {len(files)}')

        for fn in files:
            X, y = read_file(fn)
            yield X, [y]   # here I lose info about fn

    def get_shapes():
        X, _ = next(generator())
        return tf.TensorShape(X.shape), tf.TensorShape(1)

    return (
        tf.data.Dataset
            .from_generator(
                generator,
                output_types=(tf.float64, tf.uint8),
                output_shapes=get_shapes())
            .batch(128, drop_remainder=True)
            .prefetch(256))

model.fit(make_dataset(directory_train))
y_pred = model.predict(make_dataset(directory_test))

# here: what is the most elegant way to retain the input filenames for y_pred ?

Ps the shuffle in generator is needed to ensure input data randomness, since the files are being read in "lazy" mode.

Answer 1

One possibility is to make your generator return the filename, and pass that as a debug input to your model.

An end to end example.

Lets train a simple Linear Regression, here is the definition of the model:

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1,input_shape=(1,))])
model.compile(loss='mse', optimizer='sgd')

Let's generate some data

Now, lets generate some random data, and associate each sample to a dummy filename

# generating some random data
filenames = np.array([[f"{s}.txt"] for s in string.ascii_lowercase])
X = np.random.uniform(size=(26, 1))
y = np.random.uniform(size=(26, 1))
data = np.concatenate((filenames, X, y), axis=1)

Let's look a the first element of the data:

>>> data[0]
array(['a.txt', '0.36798830850651043', '0.5976948635618315'], dtype='<U32')

Training the model

We train the model with our data. For simplicity, let's use our arrays X and y .

model.fit(X,y)

Creating our debug generator and our debugging model

So, what we want is to create a debug input in our model, that is just there to receive our filename, and output it directly without any transformation. To do that, we use the Functionnal API to wrap the existing model into our debug model:

debug_input = tf.keras.Input(shape=(), dtype=tf.string)
debug_model = tf.keras.Model([debug_input, model.input], [debug_input, model.output])

And now, we need to create a dataset generator that will yield a tuple (filename, feature) to feed our debug model:

def debug_gen(data):
    # shuffling data
    np.random.shuffle(data)
    for filename, feature, label in data:
        yield (filename, feature), label


debug_ds = tf.data.Dataset.from_generator(
    lambda: debug_gen(data),
    output_types=((tf.string, tf.float64), tf.float64),
    output_shapes=((tf.TensorShape(()), tf.TensorShape(())), tf.TensorShape(())),
).batch(1)

Now, if we call predict on one item from our generator, we should get the filename as an output, as well as the prediction:

>>> debug_model.predict(debug_ds.take(1))
[array([b'a.txt'], dtype=object), array([[-0.7604195]], dtype=float32)]