Керас: модели ТПУ должны иметь постоянную форму для всех операций

Question

Я работаю с моделью предварительно подготовленного кераса и хочу запустить ее на TPU в Google Colab Laboratory, но получаю следующую ошибку:

ValueError: Слой имеет переменную форму в неразмер партии. Модели TPU должны иметь постоянные формы для всех операций.

Возможно, вам придется указать 'input_length' для слоев RNN / TimeDistributed.

Layer: Форма ввода: [(None, 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)

IsЕсть ли способ исправить размер пакета во время компиляции, чтобы избавиться от ошибки выше? Или проблема в чем-то совершенно ином?

Answer 1

Я согласен с комментариями - чтобы заставить его работать, вам нужно настроить различные переменные выходные формы (например, Нет, Нет, 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 в ведро, к которому может получить доступ колаб.

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