autokeras
autokeras copied to clipboard
keras-team
Support TPU
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
@vincent-hui Yes, we are working on this issue. Please refer to #105 . Thanks.
@haifeng-jin so is there TPU support in Colab for AutoKeras? Or do we have to stick with GPU’s?
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.
If anyone is having any error caused by running AutoKeras on TPU. Please paste it here.
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.
