bert 提取特征时TPU出错(缺少train_batch_size)

h43kikqp  于 4个月前  发布在  其他
关注(0)|答案(4)|浏览(65)

无法在Colab的TPU上运行extract_features。
当设置为使用TPU时,会出现以下错误:
!python /content/bert_repo/extract_features.py
--input_file=$input_file
--vocab_file=$VOCAB_FILE
--bert_config_file=$CONFIG_FILE
--init_checkpoint=$INIT_CHECKPOINT
--output_file=$output_file
--use_tpu=true
--batch_size=32
Traceback (most recent call last):
File "/content/bert_repo/extract_features.py", line 419, in tf.app.run()
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/app.py", line 125, in run
_sys.exit(main(argv))
File "/content/bert_repo/extract_features.py", line 382, in main
predict_batch_size=FLAGS.batch_size)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py", line 2095, in init
raise ValueError(' train_batch_size cannot be None ')
ValueError: train_batch_size cannot be None
代码正在设置predict_batch_size,但是TPU估计器似乎还需要一个train_batch_size。然而,它没有进行训练,只是进行预测,所以我不清楚为什么需要这个。

7dl7o3gd

7dl7o3gd1#

我刚遇到同样的问题。正准备打开一个问题,然后看到你6分钟前就打开了一个关于完全相同问题的。不可思议。

eimct9ow

eimct9ow2#

我遇到了同样的问题。

qzwqbdag

qzwqbdag3#

有人已经找到了解决这个问题的方法吗?

pprl5pva

pprl5pva4#

您可能需要在估算器中添加一个train_batch_size参数。

estimator = tf.contrib.tpu.TPUEstimator(
        use_tpu=FLAGS.use_tpu,
        model_fn=model_fn,
        config=run_config,
        predict_batch_size=FLAGS.batch_size,
        train_batch_size=256)

您还需要通过在命令行中添加--master="your.tpu.address"来指定TPU的地址,以便它能找到您的TPU。

相关问题