tensorflow - 将 tensorflow 数据集记录分成多条记录

Question

我有一个未批量处理的 tensorflow看起来像这样的数据集:

ds = ...
for record in ds.take(3):
    print('data shape={}'.format(record['data'].shape))

-> data shape=(512, 512, 87)
-> data shape=(512, 512, 277)
-> data shape=(512, 512, 133)

我想将数据以深度为 5 的块形式提供给我的网络。在上面的示例中，形状 (512, 512, 87) 的张量将被划分为 17 个形状 (512, 512, 5) 的张量。矩阵的最后 2 行 ( tensor[:,:, 85:87] ) 应该被丢弃。
例如:

chunked_ds = ...
for record in chunked_ds.take(1):
    print('chunked data shape={}'.format(record['data'].shape))

-> chunked data shape=(512, 512, 5)

我如何从 ds 获取至 chunked_ds ? tf.data.Dataset.window() 看起来像我需要的，但我无法让它工作。

Answer 1

这实际上可以使用 tf.data.Dataset 来完成- 仅操作:

data = tf.random.normal( shape=[ 10 , 512 , 512 , 87 ] )
ds = tf.data.Dataset.from_tensor_slices( ( data ) )
chunk_size = 5
chunked_ds = ds.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(tf.transpose(x, perm=[2, 0, 1])).batch(chunk_size, drop_remainder=True)) \
                    .map(lambda rec: tf.transpose(rec, perm=[1, 2, 0]))

那里发生了什么:
首先，我们将每条记录视为一个单独的数据集，并对其进行置换，以便最后一个维度成为批次维度( flat_map 将再次将我们的内部数据集扁平化为张量)

.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(tf.transpose(x, perm=[2, 0, 1])

然后我们按 5 分批，但我们不关心余数

.batch(chunk_size, drop_remainder=True))

最后，重新排列张量，以便我们在开始时有 512x512:

.map(lambda rec: tf.transpose(rec, perm=[1, 2, 0]))