keras-team / autokeras

AutoML library for deep learning
http://autokeras.com/
Apache License 2.0
9.1k stars 1.4k forks source link

Support TPU #269

Open vincent-hui opened 5 years ago

vincent-hui commented 5 years ago

Is it possible to support TPU?

TPU can be used in colab https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/fashion_mnist.ipynb

haifeng-jin commented 5 years ago

@vincent-hui Yes, we are working on this issue. Please refer to #105 . Thanks.

neel04 commented 3 years ago

@haifeng-jin so is there TPU support in Colab for AutoKeras? Or do we have to stick with GPU’s?

haifeng-jin commented 3 years ago

This is a really old issue. I expect we already support TPU, but having some performance issues. Should be solved when we support compile function for AutoModels.

haifeng-jin commented 3 years ago

If anyone is having any error caused by running AutoKeras on TPU. Please paste it here.

shun-lin commented 3 years ago

Looks like TPU is supported as KerasTuner now supports TPU, here is an example: https://colab.research.google.com/drive/1AaFDTQaq8B5p3-FKtNDdMdhlmUHtwOIM?authuser=1#scrollTo=HUJ0LxZnbQnc

and it's quite simple actually with only 2 lines of code:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
strategy = tf.distribute.TPUStrategy(resolver)
clf = ak.ImageClassifier(max_trials=3, distribution_strategy=strategy)

Note: It looks like using with strategy.scope() will not work (same as KerasTuner) and the TPU strategy have to be passed down this way. I think it will be great if this is added somewhere in the documentation (I'm happy to contribute if needed), also thanks for the awesome package!

Edit: Actually there are some models generate that are not compatible with TPU will break, see the same colab above.