Почему предварительно обученный кодер Swin Transformer не работает на TPU, но работает на CPU в Colab? ⇐ Python

[Anonymous]

Я работаю над задачей сегментации изображений и пытаюсь использовать предварительно обученный кодер Swin Transformer Large (Swin-L) для магистрали извлечения признаков. Код отлично работает на процессоре в Colab. Однако при переключении на 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)
 in ()
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 

Вопрос:
Почему этот код работает на ЦП, но не работает на ТПУ в Colab? Как я могу исправить эту проблему, чтобы сделать ее совместимой с выполнением TPU?
Будем очень признательны за любую информацию или рекомендации. Спасибо!

Подробнее здесь: https://stackoverflow.com/questions/792 ... -cpu-in-co