如何在 Tensorflow 中有效地加入来自 TFRecords 的数据

Question

在TFRecords 上训练 TensorFlow 模型时，我需要有效地加入少量数据。如何使用已解析的 TFRecord 中的信息进行此查找？

更多细节：

我正在使用 TFRecords 在大型数据集上训练卷积网络。每个TFRecord 都包含原始图像以及目标标签，以及有关图像的一些元数据。训练的一部分是我需要使用特定于一组图像的mean 和std 来标准化图像。为此，我将mean 和std 硬编码到TFRecord 中。然后在我的parse_example 中使用它，用于映射我的input_fn 中的Dataset，如下所示：

def parse_example(..):
    # ...
    parsed = tf.parse_single_example(value, keys_to_features)
    image_raw = tf.decode_raw(parsed['image/raw'], tf.uint16)
    image = tf.reshape(image_raw, image_shape)
    image.set_shape(image_shape)

    # pull hardcoded pixels mean and std from the parsed TFExample
    mean = parsed['mean']
    std = parsed['std']

    image = (tf.cast(image, tf.float32) - mean) / std

    # ...

    return image, label

虽然上述方法有效并且可以缩短训练时间，但它的局限性在于我经常想更改我使用的 mean 和 std。与其将mean 和std 写入TFRecords，我更愿意在训练时查找适当的汇总统计信息。这意味着当我训练时，我有一个小的 python 字典，我可以使用从TFRecord 解析的图像信息来查找适当的摘要统计信息。我遇到的问题是我似乎无法在我的张量流图中使用这个 python 字典。如果我尝试直接进行查找，它不起作用，因为我有张量对象而不是实际的原语。这是有道理的，因为 input_fn 正在为 TensorFlow 进行符号操作构建计算图（对吗？）。我该如何解决这个问题？

我尝试过的一件事是从字典中创建一个查找表，如下所示：

def create_channel_hashtable(keys, values, default_val=-1):
    initializer = tf.contrib.lookup.KeyValueTensorInitializer(keys, values)
    return tf.contrib.lookup.HashTable(initializer, default_val)

可以在parse_example 函数中创建和使用哈希表来进行查找。这一切都“有效”，但它极大地减慢了训练速度。值得注意的是，这种培训是在 TPU 上进行的。使用来自TFRecords 的值的原始方法，训练速度非常快，并且不受 IO 的限制，但是当使用哈希查找时，这种情况会发生变化。处理这些情况的建议方法是什么？虽然重新打包TFRecords 是可行的，但如果要查找的数据很小并且可以提高效率，这似乎很愚蠢。

Answer 1

这个问题涵盖了这个主题：

How can I merge multiple tfrecords file into one file?

似乎您会将 TFRecords 保存到文件中，然后使用 TFRecordDataset 将它们全部提取到一个数据集中。我链接的上述问题的答案中给出的代码是这样的：

dataset = tf.data.TFRecordDataset(filenames_to_read,
    compression_type=None,    # or 'GZIP', 'ZLIB' if compress you data.
    buffer_size=10240,        # any buffer size you want or 0 means no buffering
    num_parallel_reads=os.cpu_count()  # or 0 means sequentially reading
)