在 CloudML for GPU 上加速 TFRecords 输入 Keras 模型

Question

我想以超快的速度将 TFRecords 输入到我的模型中。然而，目前，我的 GPU（GCP 上的单 K80）的负载为 0%，这在 CloudML 上非常慢。

我在 GCS 中有 TFRecords：train_directory = gs://bucket/train/*.tfrecord，（大约 100 个大小为 30mb-800mb 的文件），但出于某种原因，它很难将数据以足够快的速度输入到我的模型中以供 GPU 使用。

有趣的是，将数据加载到内存中并使用 fit_generator() 使用 numpy 数组的速度提高了 7 倍。在那里我可以指定多处理和多工人。

我当前的设置解析 tf 记录并加载无限 tf.Dataset。理想情况下，该解决方案将在内存中保存/完善一些批次，供 GPU 按需使用。

def _parse_func(record):
    """ Parses TF Record"""
    keys_to_features = {}
    for _ in feature_list: # 300 features ['height', 'weights', 'salary']
         keys_to_features[_] = tf.FixedLenFeature([TIME_STEPS], tf.float32)
    parsed = tf.parse_single_example(record, keys_to_features)
    t = [tf.manip.reshape(parsed[_], [-1, 1]) for _ in feature_list]
    numeric_tensor = tf.concat(values=t, axis=1)

    x = dict()
    x['numeric'] = numeric_tensor
    y = ...
    w = ...

    return x, y, w

def input_fn(file_pattern, b=BATCH_SIZE):
    """
    :param file_pattern: GCS bucket to read from
    :param b: Batch size, defaults to BATCH_SIZE in hparams.py
    :return: And infinitely iterable data set using tf records of tf.data.Dataset class
    """
    files = tf.data.Dataset.list_files(file_pattern=file_pattern)
    d = files.apply(
        tf.data.experimental.parallel_interleave(
            lambda filename: tf.data.TFRecordDataset(filename),
            cycle_length=4,
            block_length=16,
            buffer_output_elements=16,
            prefetch_input_elements=16,
            sloppy=True))
    d = d.apply(tf.contrib.data.map_and_batch(
        map_func=_parse_func, batch_size=b,
        num_parallel_batches=4))
    d = d.cache()
    d = d.repeat()
    d = d.prefetch(1)
    return d

获取火车数据

# get files from GCS bucket and load them into dataset
train_data = input_fn(train_directory, b=BATCH_SIZE)

拟合模型

model.fit(x=train_data.make_one_shot_iterator())

我在 CloudML 上运行它，所以 GCS 和 CloudML 应该很快。

---

CloudML CPU 使用率：

正如我们在下面看到的，CPU 为 70%，内存没有增加超过 10%。那么dataset.cache() 有什么作用呢？

CloudML 日志中的 GPU 指标

如下图所示，GPU 似乎已关闭！内存也为0mb。缓存存放在哪里？

没有在 GPU 上运行的进程！

编辑：

似乎确实没有在 GPU 上运行进程。我试图明确说明：

tf.keras.backend.set_session(tf.Session(config=tf.ConfigProto(
    allow_soft_placement=True,
    log_device_placement=True)))

train_data = input_fn(file_pattern=train_directory, b=BATCH_SIZE)

model = create_model()

with tf.device('/gpu:0'):
    model.fit(x=train_data.make_one_shot_iterator(),
              epochs=EPOCHS,
              steps_per_epoch=STEPS_PER_EPOCH,
              validation_data=test_data.make_one_shot_iterator(),
              validation_steps=VALIDATION_STEPS)

但一切仍然使​​用 CPU！

Answer 1

在我的例子中，我使用了一个自定义的 setup.py 文件，它使用了一个仅 CPU 的 Tensorflow 版本。

我在踢自己，请安装tensorflow-gpu。