首页 > 解决方案 > 获取BERT最后4层的输出(Tensorflow)

问题描述

我正在阅读 Andreas Poyiatzis 发表在向数据科学网站上的文章。下面是文章的链接。

https://towardsdatascience.com/nlp-extract-contextualized-word-embeddings-from-bert-keras-tf-67ef29f60a7b#--responses

生成的嵌入是通过使用 TPU。但我想在不支持 TPU 的本地机器上运行它。有人可以让我知道如何在 CPU 上运行相同的代码吗?您的帮助将不胜感激。谢谢!

标签: tensorflowtensorflow2.0bert-language-modeltputensorflow-hub

解决方案


在链接之后,代码说如果没有可用的 TPU,Estimator 会退回到 CPU 或 GPU。只要确保你use_tpu = False在它出现的任何地方设置:

  # If TPU is not available, this will fall back to normal Estimator on CPU
  # or GPU.
  estimator = tf.contrib.tpu.TPUEstimator(
      use_tpu=False,
      model_fn=model_fn,
      config=run_config,
      predict_batch_size=BATCH_SIZE,
      train_batch_size=BATCH_SIZE)

也直接在model_fn_builder调用中。像这样:

model_fn = model_fn_builder(
      bert_config=bert_config,
      init_checkpoint=INIT_CHECKPOINT,
      layer_indexes=layer_indexes,
      use_tpu=False,
      use_one_hot_embeddings=True)

推荐阅读