Comment enregistrer un ensemble de données Tensorflow dans un fichier ?

TFRecordWriter semble être l'option la plus pratique, mais malheureusement, il ne peut écrire des ensembles de données qu'avec un seul tenseur par élément. Voici quelques solutions de contournement que vous pouvez utiliser. Tout d'abord, puisque tous vos tenseurs ont le même type et une forme similaire, vous pouvez tous les concaténer en un seul et les diviser ultérieurement lors du chargement :

import tensorflow as tf

# Write
a = tf.zeros((100, 512), tf.int32)
ds = tf.data.Dataset.from_tensor_slices((a, a, a, a[:, 0]))
print(ds)
# <TensorSliceDataset shapes: ((512,), (512,), (512,), ()), types: (tf.int32, tf.int32, tf.int32, tf.int32)>
def write_map_fn(x1, x2, x3, x4):
    return tf.io.serialize_tensor(tf.concat([x1, x2, x3, tf.expand_dims(x4, -1)], -1))
ds = ds.map(write_map_fn)
writer = tf.data.experimental.TFRecordWriter('mydata.tfrecord')
writer.write(ds)

# Read
def read_map_fn(x):
    xp = tf.io.parse_tensor(x, tf.int32)
    # Optionally set shape
    xp.set_shape([1537])  # Do `xp.set_shape([None, 1537])` if using batches
    # Use `x[:, :512], ...` if using batches
    return xp[:512], xp[512:1024], xp[1024:1536], xp[-1]
ds = tf.data.TFRecordDataset('mydata.tfrecord').map(read_map_fn)
print(ds)
# <MapDataset shapes: ((512,), (512,), (512,), ()), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

Mais, plus généralement, vous pouvez simplement avoir un fichier séparé par tenseur et ensuite tous les lire :

import tensorflow as tf

# Write
a = tf.zeros((100, 512), tf.int32)
ds = tf.data.Dataset.from_tensor_slices((a, a, a, a[:, 0]))
for i, _ in enumerate(ds.element_spec):
    ds_i = ds.map(lambda *args: args[i]).map(tf.io.serialize_tensor)
    writer = tf.data.experimental.TFRecordWriter(f'mydata.{i}.tfrecord')
    writer.write(ds_i)

# Read
NUM_PARTS = 4
parts = []
def read_map_fn(x):
    return tf.io.parse_tensor(x, tf.int32)
for i in range(NUM_PARTS):
    parts.append(tf.data.TFRecordDataset(f'mydata.{i}.tfrecord').map(read_map_fn))
ds = tf.data.Dataset.zip(tuple(parts))
print(ds)
# <ZipDataset shapes: (<unknown>, <unknown>, <unknown>, <unknown>), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

Il est possible d'avoir l'ensemble de données dans un seul fichier avec plusieurs tenseurs séparés par élément, à savoir sous la forme d'un fichier de TFRecords contenant tf.train.Example s, mais je ne sais pas s'il existe un moyen de les créer dans TensorFlow, c'est-à-dire sans avoir à extraire les données de l'ensemble de données dans Python, puis à les écrire dans le fichier d'enregistrements.

Un incident était ouvert sur GitHUb et il semble qu'il y ait une nouvelle fonctionnalité disponible dans TF 2.3 pour écrire sur le disque :

https://www.tensorflow.org/api_docs/python/tf/data/experimental/savehttps://www.tensorflow.org/api_docs/python/tf/data/experimental/load

Je n'ai pas encore testé cette fonctionnalité, mais elle semble faire ce que vous voulez.