Keras：TPU模型的所有操作都必须具有恒定的形状

Question

我正在使用预训练的keras模型，并且想通过Google Colaboratory在TPU上运行它，但是出现以下错误：

ValueError：图层在非批量尺寸中具有可变的形状。对于所有操作，TPU模型必须具有恒定的形状。

您可能必须为RNN / TimeDistributed图层指定“ input_length”。

图层：输入形状：[（无，128，768），（无，1）]输出形状：（无，无，768）

我正在使用keras-xlnet。据我了解，按此处和此处说明编译模型时，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)

有什么办法可以解决我在编译时的批处理大小，以摆脱上面的错误？还是问题完全不同？

Answer 1

我同意这些评论 - 要使其正常工作，您需要将各种可变输出形状（例如 None、None、768）调整为固定大小（第一个批次尺寸除外）。也许你可以通过简单的填充来做到这一点。如果您可以循环遍历已保存的模型层并将权重加载到您用填充尺寸编写的新模型中，那么它甚至可能有效。我想说，这比考虑 TPU 就绪版本已经可用更麻烦。

我建议在此模型中放弃 Keras。官方 TensorFlow XLNet 实现应该无需修改即可与 TPU 一起使用。它还配备了预先训练的检查点。https://github.com/zihangdai/xlnet

它使用标准 TPUEstimator 类将模型函数发送到 TPU 工作线程，因此您无需费力处理tf.contrib.tpu.keras_to_tpu_model.

存储库中给出的示例可以在 colab 中运行$TPU_NAME，$COLAB_TPU_ADDR您可以将预训练的检查点和 imdb 数据上传到 colab 可以访问的存储桶。

python run_classifier.py \
  --use_tpu=True \
  --tpu=${TPU_NAME} \
  --do_train=True \
  --do_eval=True \
  --eval_all_ckpt=True \
  --task_name=imdb \
  --data_dir=${IMDB_DIR} \
  --output_dir=${GS_ROOT}/proc_data/imdb \
  --model_dir=${GS_ROOT}/exp/imdb \
  --uncased=False \
  --spiece_model_file=${LARGE_DIR}/spiece.model \
  --model_config_path=${GS_ROOT}/${LARGE_DIR}/model_config.json \
  --init_checkpoint=${GS_ROOT}/${LARGE_DIR}/xlnet_model.ckpt \
  --max_seq_length=512 \
  --train_batch_size=32 \
  --eval_batch_size=8 \
  --num_hosts=1 \
  --num_core_per_host=8 \
  --learning_rate=2e-5 \
  --train_steps=4000 \
  --warmup_steps=500 \
  --save_steps=500 \
  --iterations=500