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

Question

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

当我启动 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，因为在available TPU operations 列表中没有preprocessing 选项。我说的对吗？

推理时有多个Benefits of doing preprocessing inside the model。

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

谢谢。

大卫

Answer 1

你说对了一半。 TPU operations 的列表包括较低级别的 TF 函数，但不包括 Keras 层。从错误消息来看，您的预处理层似乎试图在图中实例化 ImageProjectiveTransformV2 操作，这是不受支持的。

作为兼容 TPU 的替代品，我建议您查看 TF 模型园中的 official EfficientNet implementation。特别是preprocessing.py 可能对您有帮助。

Answer 2

一种可能的解决方法是将图层合并到输入管道中。这有点骇人听闻，但我已经对其进行了简短的测试，它似乎可以在 TPU 上运行。例如，如果您正在使用tf.data.Dataset API，您可以创建一个图层对象，然后在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 正式解决之前，我打算将此作为一种解决方法。