Lasertagger Colab Notebook request

[daltonj]

The goal is to run lasertagger in a Google Colab notebook, similar to the BERT finetuning notebooks.

A few issues involved:

• Requirements -- Install BERT
• Colab works only with TF 1.x for TPU support
• Default 256 batch size results in OOM error on colab k80 GPU
• Add support for training / export using TFHub / Google Cloud.

[daltonj]

Start on a colab notebook
https://colab.research.google.com/drive/1G2-EPTbQMQ_QWDh5Ht-p38jFBLc1SUyi

[varepsilon]

Looks like one needs to store the data on Google Cloud bucket in order to be able to use Google's TPU.

E.g., for the BERT data:*

INPUT_DIR = 'gs://{}/bert/{}'.format(BUCKET, BERT_MODEL)
tf.gfile.MakeDirs(INPUT_DIR)

for f in tf.gfile.Glob('/content/{}/*'.format(BERT_MODEL)):
  tf.gfile.Copy(f, os.path.join(INPUT_DIR, f.split('/')[-1]))

%env BERT_BASE_DIR=$INPUT_DIR

And for the output data:

GS_OUTPUT_DIR = 'gs://{}/output'.format(BUCKET)
tf.gfile.MakeDirs(GS_OUTPUT_DIR)

for f in tf.gfile.Glob('/content/output/*'):
  if tf.gfile.Stat(f).is_directory:
    continue
  tf.gfile.Copy(f, os.path.join(GS_OUTPUT_DIR, f.split('/')[-1]))

%env OUTPUT_DIR=$GS_OUTPUT_DIR

(*) Assuming the data was previously stored locally, e.g.,

bert_url = 'https://storage.googleapis.com/bert_models/2018_10_18/' + BERT_MODEL + '.zip'
bert_zip = BERT_MODEL + '.zip'
!wget $bert_url
!unzip $bert_zip

[daltonj]

The notebook has been updated to support TPU, via writing all data to a cloud bucket.

Note: the prediction step - predict_main.py doesn't have flags for using a TPU. It's currently running very slowly, roughly 1 example per second. Is this expected?

I1107 15:31:14.475278 139814862706560 predict_main.py:89] 0 examples processed, 0 converted to tf.Example.
I1107 15:33:06.323075 139814862706560 predict_main.py:89] 100 examples processed, 100 converted to tf.Example.

[ekQ]

Hi and sorry for the slow reply! Having a Colab would indeed be very useful but atm I don't have time to do that. However, if you'd like to create a pull request, I'd be very happy to review it.

Regarding slow inference: This is indeed an issue and the expected behavior when you run the code as such. Internally, we heavily parallelize inference so it's not an issue in that case. To make it faster, one should ideally increase the batch size (currently it's 1 [*]) which requires small code changes.

A quicker fix is to use LaserTaggerFF by setting use_t2t_decoder to false in configs/lasertagger_config.json. This should already make prediction about 40 times faster (at least on GPU). This may hurt the accuracy slightly but not radically at least in our experiments.

[*] https://github.com/google-research/lasertagger/blob/master/predict_utils.py#L57