Running BicleanerAI with TPUs

[charliekocsis]

Hi Bitextor team,

I'm using Bicleaner AI to clean up Opus corpora. As part of that I was looking at using TPUs for the clean up process.

On Kaggle, there is documentation on using TPUs with TF and Keras as follows:

# detect and init the TPU
tpu = tf.distribute.cluster_resolver.TPUClusterResolver()

# instantiate a distribution strategy
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = tf.distribute.TPUStrategy(tpu)

# instantiating the model in the strategy scope creates the model on the TPU
with tpu_strategy.scope():
    model = tf.keras.Sequential( … ) # define your model normally
    model.compile( … )

# train model normally
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=…)

But I am not sure how I would compile and run a bicleaner-ai model with TPUs. Is there a simple way to integrate Bicleaner AI with TPU support?

Thanks for your help in advance.

[ZJaume]

I guess that if you need to set the TPU scope during model loading (not exactly as the example, because that is for training, though), you will need to do it here: https://github.com/bitextor/bicleaner-ai/blob/0547fa3b68838edecbb889aa50d22e602eb36c42/src/bicleaner_ai/models.py#L225-L230

But never tried using TPUs with Bicleaner AI, not sure if it will work.