tensorflow - Colab：在 TPU 上缓存数据集

Question

我想设置类似于以下 colab notebook 的东西。我有一个 100MB 的单 TFRecord，我想用 TPU 训练它。

我的训练输入函数如下：

def train_input_fn(batch_size=1024):
  dataset = tf.data.TFRecordDataset(TRAIN_RECORD)
  dataset = dataset.cache()
  dataset = dataset.repeat()
  dataset = dataset.shuffle(100)
  dataset = dataset.map(parse_fn)
  dataset = dataset.batch(batch_size, drop_remainder=True)
  return dataset

据我了解，使用 TPU 时，数据集不能驻留在机器硬盘上，这就是我添加dataset.cache(). 但我仍然得到

UnimplementedError（参见上面的回溯）：文​​件系统方案“[本地]”未实现（文件：“train.tfrecord”）

Answer 1

TPU 需要谷歌云存储。不支持本地存储。

https://cloud.google.com/tpu/docs/troubleshooting#cannot_use_local_filesystem

Answer 2

看起来错误出现在dataset = tf.data.TFRecordDataset(TRAIN_RECORD)与本地 fs 一起使用的这一行上。我认为您应该像在笔记本中那样在训练功能之外加载数据。