将 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) 的张量。tensor[:,:, 85:87]应丢弃矩阵 ( )的最后 2 行。

例如：

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这实际上可以使用-only 操作来完成：

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]))