【发布时间】：2020-02-25 04:08:12
【问题描述】：
我正在使用预训练的 keras 模型，我想在 Google Colaboratory 的 TPU 上运行它，但我收到以下错误：
值错误： 层在非批量维度上具有可变形状。 TPU 型号必须 所有操作的形状都保持不变。
您可能必须为 RNN/TimeDistributed 层指定“input_length”。
层： 输入形状：[(None, 128, 768), (None, 1)] 输出形状：（无，无，768）
我正在与keras-xlnet 合作。据我了解，按照here 和here 的说明编译模型时，TPU 需要具有固定的批量大小。
模型从检查点加载：
from keras_xlnet import Tokenizer, load_trained_model_from_checkpoint,
ATTENTION_TYPE_BI
checkpoint_path = 'xlnet_cased_L-12_H-768_A-12'
tokenizer = Tokenizer(os.path.join(checkpoint_path, 'spiece.model'))
model = load_trained_model_from_checkpoint(
config_path=os.path.join(checkpoint_path, 'xlnet_config.json'),
checkpoint_path=os.path.join(checkpoint_path, 'xlnet_model.ckpt'),
batch_size=BATCH_SIZE,
memory_len=512,
target_len=SEQ_LEN,
in_train_phase=False,
attention_type=ATTENTION_TYPE_BI,
)
model.summary()
然后编译模型（经过一些更改）：
from keras_bert import AdamWarmup, calc_train_steps
decay_steps, warmup_steps = calc_train_steps(
y_train.shape[0],
batch_size=BATCH_SIZE,
epochs=EPOCHS,
)
model.compile(
AdamWarmup(decay_steps=decay_steps, warmup_steps=warmup_steps, lr=LR),
loss='binary_crossentropy',
)
然后，模型被加载到 TPU，出现错误：
tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
strategy = tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
)
with tf.keras.utils.custom_object_scope(get_custom_objects()):
tpu_model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)
有没有办法可以在编译时修复批量大小以消除上述错误？还是问题完全不同？
【问题讨论】：
-
请注意，在 TPU 上运行可能很难解决所有各种错误，主要是 TensorFlow 自己过时的模块；即使你解决了这个问题，很可能还会有另一个问题。此外，功能主要限于核心 TF 模块（而不是例如自定义优化器）。最后，TPU 的速度非常快，但除非你的 input data pipeline 非常高效，否则你不会看到太大的收益；数据加载时间很容易成为瓶颈。
-
感谢您的评论。我认为我的输入管道等运行良好，因为它主要与我在keras-bert 中使用的设置相同。 GPU 在我的训练集（使用 BERT）上要慢很多，所以在整个数据集上运行它没有意义。
-
还好，TPU应该值得调试
-
恐怕你需要自己重写模型，固定形状并复制权重。但是像在这个问题中添加
model.summary()的信息（包括子模型，如果有的话）肯定会有助于检测更明显的东西。 -
模型摘要已找到here，因为它太大，无法在此处有用地上传。
标签： tensorflow keras google-colaboratory tpu google-cloud-tpu