如何在云 TPU 上使用 Tensorflow Hub 模型？

Question

我正在尝试使用来自 Kaggle 上的 tensorflow hub 的模型。像这样：

m = tf.keras.Sequential([
hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4", output_shape=[1280],
               trainable=False),  # Can be True, see below.
tf.keras.layers.Dense(num_classes, activation='softmax')
])

m.build([None, 224, 224, 3])  # Batch input shape.

它在 GPU 上运行良好，但是一旦我切换到带有 TF 记录的 TPU，我就会收到以下错误：

InvalidArgumentError: Unsuccessful TensorSliceReader constructor: Failed to get matching files on /tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables: Unimplemented: File system scheme '[local]' not implemented (file: '/tmp/tfhub_modules/87fb99f72aec02d017e12c0a3d86c5c182ec22ca/variables/variables')

然而，设置和 tfrecords 数据集都是正确的，因为它可以将预训练模型切换到相同模型的 keras 应用程序（例如，上面使用 mobilenet keras 应用程序）。

我试过缓存但没有成功，在遵循本指南时有什么我必须注意的吗： https://www.tensorflow.org/hub/caching

提前致谢！

Answer 1

发生失败是因为 TPU 正在尝试从它无权访问的 /tmp/ 加载 TFHub 模型。您应该能够通过以下方式使其正常工作：

with strategy.scope():
  load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')
  m = tf.keras.Sequential([
    hub.KerasLayer(
      "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4",
      output_shape=[1280],
      load_options=load_locally,
      trainable=False),  # Can be True, see below.
    tf.keras.layers.Dense(num_classes, activation='softmax')
  ])

来源：EfficientNetB7 on 100+ flowers.

如何在云 TPU 上使用 Tensorflow Hub 模型？

How to use Tensorflow Hub models on cloud TPU?

keras

tensorflow

tpu