tensorflow - Keras:TPU 模型的所有操作必须具有恒定的形状
问题描述
我正在使用预训练的 keras 模型,我想通过 Google Colaboratory 在 TPU 上运行它,但出现以下错误:
ValueError:图层在非批量维度中具有可变形状。TPU 模型对于所有操作必须具有恒定的形状。
您可能必须为 RNN/TimeDistributed 层指定“input_length”。
层:输入形状:[(None, 128, 768), (None, 1)] 输出形状:(None, None, 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)
有没有办法可以在编译时修复批量大小以消除上述错误?还是问题完全不同?
解决方案
我同意这些评论——要让它工作,您需要将各种可变输出形状(例如无、无、768)调整为固定大小(第一批尺寸除外)。也许你可以用简单的填充来做到这一点。如果您可以遍历保存的模型层并将权重加载到您使用填充尺寸编写的新模型中,它甚至可以工作。我想说这比考虑 TPU 就绪版本更麻烦。
我建议这个模型远离 Keras。官方的 TensorFlow XLNet 实现无需修改即可与 TPU 一起使用。它还带有预训练的检查点。https://github.com/zihangdai/xlnet
它使用标准的 TPUEstimator 类将模型函数发送给 TPU 工作人员,因此您无需乱用tf.contrib.tpu.keras_to_tpu_model
.
存储库中给出的示例可以在 colab where $TPU_NAME
is中运行,$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
推荐阅读
- microsoft-graph-toolkit - 特定 planID 的任务并分配给我
- github - 在公共网站上使用 Github 个人访问令牌
- centos7 - 拒绝访问的原因“制作缓存目录时出错?”
- formik - 使用 Yup 对多个字段进行条件验证
- django - 需要从 Django 中的数据库中获取特定的外键对象产品
- jointjs - JointJs:在节点周围的任何地方连接线/链接
- html - 为什么粘性位置在我的代码中不起作用?
- azure - AzureML 在管道之间传递数据而不保存它
- fabricjs - 使用动态文本(如 canva)创建自定义视频
- c# - C# - 使用 .Net Core 5.0 进行端口映射