Running on TPUs¶
All Tensorflow-compatible rerankers support training and inference on Google TPUs. Capreolus has been tested with both v2-8 TPUs and v3-8 TPUs.
TPUs stream their data from buckets in Google Cloud Storage rather than reading their input data from the local machine (i.e., the machine running Capreolus). Capreolus will automatically preprocess and upload the data to this bucket. However, note that GCS is not free and the user is responsible for manually deleting this data once it is no longer needed.
To use a TPU with a Tensorflow-compatible
Reranker (i.e., a reranker that depends on the
tensorflow Trainer module), set the following config options:
tpuname: the name of your TPU, such as mytpu1
tpuzone: the cloud zone your TPU is in, such as us-central1-f
storage: path to a GCS bucket where data should be stored, such as gs://your-bucket/abc/
usecache=Truewith the trainer and extractor
After setting these options, you can run Capreolus as normal. Watch for INFO logging messages at the beginning of training to confirm the TPU is being used.
While any Tensorflow-compatible Reranker can be used with TPUs, this will actually slow down small models like KNRM. TPUs are most useful with large Transformer-based models.
The following models are good candidates for running on TPUs:
TFBERTMaxP(config=None, provide=None, share_dependency_objects=False, build=True)
TensorFlow implementation of BERT-MaxP.
Deeper Text Understanding for IR with Contextual Neural Language Modeling. Zhuyun Dai and Jamie Callan. SIGIR 2019. https://arxiv.org/pdf/1905.09217.pdf
TFCEDRKNRM(config=None, provide=None, share_dependency_objects=False, build=True)
TensorFlow implementation of CEDR-KNRM. Equivalant to BERT-KNRM when cls=None.
CEDR: Contextualized Embeddings for Document Ranking Sean MacAvaney, Andrew Yates, Arman Cohan, and Nazli Goharian. SIGIR 2019. https://arxiv.org/pdf/1904.07094
TFParade(config=None, provide=None, share_dependency_objects=False, build=True)
TensorFlow implementation of PARADE.
PARADE: Passage Representation Aggregation for Document Reranking. Canjia Li, Andrew Yates, Sean MacAvaney, Ben He, and Yingfei Sun. arXiv 2020. https://arxiv.org/pdf/2008.09093.pdf