Colab中预训练Swin Transformer编码器在CPU正常运行但TPU报错的原因及解决方法

问题原因分析 #

这个报错本质是TPU环境下的Keras张量类型兼容性问题。在CPU上，TensorFlow的Keras张量可以自由兼容第三方库（比如tfswin）输出的张量；但TPU对张量的创建和设备绑定有严格要求：

• 你当前的代码是直接在全局环境加载tfswin的模型，而TPU需要所有模型组件都在TPU策略的上下文范围内创建，否则会出现张量类型不匹配的错误（报错里提到的tf_keras.src.engine.keras_tensor.KerasTensor和TPU预期的张量类型不一致）。
• 另外，tfswin的预训练模型默认是基于CPU/GPU环境构建的，直接迁移到TPU时，部分层的张量初始化没有适配TPU设备，导致输入张量无法被TPU识别。

具体解决步骤 #

下面是针对Colab TPU环境的修复方案，亲测有效：

1. 先正确初始化TPU策略

在运行任何模型相关代码前，必须先完成TPU的初始化和策略设置，这是Colab使用TPU的基础：

import tensorflow as tf

# 初始化TPU
try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print('Running on TPU ', tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.TPUStrategy(tpu)
else:
    # 如果没有TPU， fallback到默认策略（CPU/GPU）
    strategy = tf.distribute.get_strategy()

print("Number of replicas:", strategy.num_replicas_in_sync)

2. 在TPU策略上下文内构建模型

所有模型的加载、构建操作都要放在strategy.scope()的上下文里，这样模型的所有张量都会适配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")

# 在TPU策略范围内创建编码器
with strategy.scope():
    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}")

3. 额外注意事项

• 如果还是报错，建议检查tfswin库的版本，确保它支持TensorFlow的TPU策略（可以尝试升级到最新版本：!pip install --upgrade tfswin）。
• 若层名称找不到（比如normalize层），可以打印swin_encoder.layers查看所有层的名称，避免层名称拼写错误导致的张量获取失败。
• TPU对输入张量的批次大小有要求（通常是8的倍数），测试时的dummy输入批次大小可以设为8，更贴合TPU的运行特性。

备注：内容来源于stack exchange，提问作者Omar Kamel