bicleaner-ai icon indicating copy to clipboard operation
bicleaner-ai copied to clipboard

Running BicleanerAI with TPUs

Open charliekocsis opened this issue 1 year ago • 1 comments

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.

charliekocsis avatar May 22 '24 22:05 charliekocsis

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.

ZJaume avatar May 23 '24 09:05 ZJaume

Closing this for now, feel free to re-open it if there are more doubts.

ZJaume avatar Oct 31 '24 10:10 ZJaume