Colab + TPU не поддерживает TF 2.3.0 tf.keras.layers.experimental.preprocessing

Question

Я обновлял свою модель, используя TF 2.3.0 на Colab + TPU на основе https://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. Я прав?

Есть несколько преимуществ выполнения предварительной обработки внутри модели во время вывода.

Где я могу найти возможную дату реализации?

Спасибо.

Давиде

Answer 1

Возможный обходной путь - включить слои во входной конвейер. Это немного похоже на взлом, но я кратко его протестировал, и, похоже, он работает на 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)

Это не компилирует Уровни дополнения в модель, как изначально предполагалось, но это должно позволить вам увеличивать ваши обучающие данные, не требуя стороннего плагина. Я намерен использовать это как обходной путь, пока проблема не будет официально решена.

Answer 2

Вы наполовину правы. Список операций TPU включает функции TF нижнего уровня, но не слои Keras. Из сообщения об ошибке видно, что ваш уровень предварительной обработки пытается создать экземпляр операции ImageProjectiveTransformV2 на графике, что не поддерживается.

В качестве альтернативы, совместимой с TPU, я рекомендую вам взглянуть на официальный Efficien tNet реализация в саду модели TF. В частности, preprocessing.py может быть вам полезен.