为什么预训练的 Swin Transformer 编码器在 TPU 上失败，但在 Colab 的 CPU 上正常？

Question

我正在研究图像分割任务，并尝试使用预训练的 Swin Transformer Large (Swin-L) 编码器作为特征提取主干。该代码在 Colab 的 CPU 上完美运行。但是，当切换到 TPU 时，它会抛出如下所示的错误。

代码：

from tensorflow.keras import layers, Model, Input
from tfswin import SwinTransformerLarge224


def load_swin_encoder(input_shape=(512, 512, 3)):
    # Load pre-trained Swin-L model
    swin_encoder = SwinTransformerLarge224(include_top=False, weights='imagenet',
                                           input_shape=input_shape)

    # Freeze the pre-trained layers
    for layer in swin_encoder.layers:
        layer.trainable = False

    # Extract outputs from the four stages
    stage_outputs = [
        swin_encoder.get_layer('normalize').output,  # Output from the 0 stage
        swin_encoder.get_layer('layers.0').output,   # Output from the first stage
        swin_encoder.get_layer('layers.1').output,   # Output from the second stage
        swin_encoder.get_layer('layers.2').output,   # Output from the third stage
        swin_encoder.get_layer('layers.3').output,   # Output from the fourth stage
    ]
    return Model(swin_encoder.input, stage_outputs, name="SwinTransformerEncoder")


# Test Code
encoder = load_swin_encoder(input_shape=(512, 512, 3))
dummy_input = tf.random.uniform((1, 512, 512, 3))
encoder_outputs = encoder(dummy_input)

for i, output in enumerate(encoder_outputs):
    print(f"Stage {i + 1} output shape: {output.shape}")

错误：

代码在 TPU 上抛出以下错误：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-28-3cb122d32678> in <cell line: 2>()
      1 # loading Sanity check
----> 2 encoder = load_swin_encoder(input_shape=(512, 512, 3))
      3 dummy_input = tf.random.uniform((1, 512, 512, 3))
      4 encoder_outputs = encoder(dummy_input)
      5 

2 frames
/usr/local/lib/python3.10/dist-packages/keras/src/models/functional.py in __init__(self, inputs, outputs, name, **kwargs)
    117         for x in flat_inputs:
    118             if not isinstance(x, backend.KerasTensor):
--> 119                 raise ValueError(
    120                     "All `inputs` values must be KerasTensors. Received: "
    121                     f"inputs={inputs} including invalid value {x} of "

ValueError: All `inputs` values must be KerasTensors. Received: inputs=KerasTensor(type_spec=TensorSpec(shape=(None, 512, 512, 3), dtype=tf.float32, name='input_4'), name='input_4', description="created by layer 'input_4'") including invalid value KerasTensor(type_spec=TensorSpec(shape=(None, 512, 512, 3), dtype=tf.float32, name='input_4'), name='input_4', description="created by layer 'input_4'") of type <class 'tf_keras.src.engine.keras_tensor.KerasTensor'>

问题：

为什么此代码在 Colab 中的 CPU 上可以运行，但在 TPU 上却失败？如何解决此问题以使其与 TPU 执行兼容？

任何见解或指导将不胜感激。谢谢！

Answer 1

这看起来不像硬件加速器 - GPU/TPU 错误，因为我尝试通过选择可用的 - CPU/GPU 和 TPU 使用 Google Colab 复制相同的错误，但没有注意到任何问题。请找到相同内容的复制要点。

请安装最新的

Tensorflow 2.18

和

Keras 3.7

版本重试，如果问题仍然存在，请告诉我们。谢谢。