Colab+TPU 不支持 TF 2.3.0 tf.keras.layers.experimental.preprocessing

我正在使用更新我的模型TF 2.3.0基于Colab+TPUhttps://keras.io/examples/vision/image_classification_efficientnet_fine_tuning/，特别是在数据增强和从预训练权重进行迁移学习段落之后。

当我启动时model.fit我收到这个错误：

InvalidArgumentError: 9 root error(s) found.
  (0) Invalid argument: {{function_node __inference_train_function_372657}} Compilation failure: Detected unsupported operations when trying to compile graph cluster_train_function_12053586239504196919[] on XLA_TPU_JIT: ImageProjectiveTransformV2 (No registered 'ImageProjectiveTransformV2' OpKernel for XLA_TPU_JIT devices compatible with node {{node EfficientNet/img_augmentation/random_rotation_2/transform/ImageProjectiveTransformV2}}){{node EfficientNet/img_augmentation/random_rotation_2/transform/ImageProjectiveTransformV2}}
    TPU compilation failed
     [[tpu_compile_succeeded_assert/_6138790737589773377/_7]]
     [[TPUReplicate/_compile/_14198390524791994190/_6/_238]]
 

我估计TPU还不支持tf.keras.layers.experimental.preprocessing因为在列表中可用的 TPU 操作没有preprocessing选项。我对吗？

有多个在模型内部进行预处理的好处在推理时。

我在哪里可以找到可能的实施日期？

Thanks.

Davide

一种可能的解决方法是将这些层合并到输入管道中。这有点像 hack，但我对其进行了简单测试，它似乎可以在 TPU 上运行。例如，如果您正在使用tf.data.DatasetAPI，您可以创建一个图层对象，然后在其中调用它Dataset.map()将增强应用到管道：

# dummy data
images = tf.random.uniform((10, 224, 224, 1))
labels = tf.zeros((10, 1))
ds = tf.data.Dataset.from_tensor_slices((images, labels))
ds = ds.batch(10)

# now incorporate the augmentation 'layer' into the pipeline
augmentor = tf.keras.layers.experimental.preprocessing.RandomRotation((-0.1, 0.1))
# augment the images, pass the labels through untouched
ds = ds.map(lambda x, y: (augmentor.call(x), y))

# assume we've compiled a model elsewhere
model.fit(ds)

这不会按照最初的预期将增强层编译到模型中，但它应该允许您在不需要第三方插件的情况下增强训练数据。我打算用这个作为解决方法，直到issue已正式解决。