Bert: extract_features on TPU error (missing train_batch_size)

[daltonj]

I am unable to run extract_features on a TPU in colab.
https://colab.research.google.com/drive/1tMi3GYv74IA5N4oB_60BylUOQaVJdDkT

When set to use the TPU, it gives the following error:

!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

The code is setting the predict_batch_size, but the TPU estimator also appears to need a train_batch_size. However, it's not doing training, just prediction, so it's unclear to me why this should be required.

[ivsanro1]

I just encountered the same issue. Was going to open an issue and I see you opened one about the exact same problem just 6 minutes ago. Incredible.

[NguyenAnhTien]

I have got the same issue

[blkfin]

Anyone have a workaround for this yet?

[chenqinkai]

You may need to add a train_batch_size parameter to estimator

    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)

You also need to specify the address of your TPU by adding --master="your.tpu.address" in the command line to allow it to find your TPU.