GPU or TPU not detected with GPU/TPU runtime post mediapipe & mediapipe_model_maker installation

[PranayLohia]

Have I written custom code (as opposed to using a stock example script provided in MediaPipe)

Yes

OS Platform and Distribution

MacOS & Linux

Python Version

3.10.12

MediaPipe Model Maker version

0.2.1.4

Task name (e.g. Image classification, Gesture recognition etc.)

Model Training

Describe the actual behavior

I am unable to run the MobileBERT model training on GPU & TPU with proper GPU & TPU runtime, it pick tf.config.list_physical_devices() only as CPU

Describe the expected behaviour

Expected to run on all accelerators

Standalone code/steps you may have used to try to get what you need

-- coding: utf-8 -- """training_gpu.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/10GHWGPAZJKM8RbXCpbBt2mit_gs3txbn?resourcekey=0-f-FpXhUgIbZbFHyz0sEC0g """ # @title Pre Import tensorflows import tensorflow as tf assert tf.version.startswith('2') tf.get_logger().setLevel('ERROR') import warnings warnings.filterwarnings('ignore') warnings.filterwarnings('ignore', category=UserWarning) tf.config.list_physical_devices() # @title Add Imports !pip install mediapipe-model-maker # @title Import tensorflows import tensorflow as tf assert tf.version.startswith('2') tf.get_logger().setLevel('ERROR') import warnings warnings.filterwarnings('ignore') warnings.filterwarnings('ignore', category=UserWarning) tf.config.list_physical_devices() from mediapipe_model_maker import text_classifier

Other info / Complete Logs

[kuaashish]

Hi @PranayLohia,

Colab example you provided requires access. Could you please share it again so we can better understand the issue?

Thank you!!

[PranayLohia]

provided access to u https://colab.research.google.com/drive/1WD-s6Ozo0w7Y5Z6fZZCiWIDrJf8LOmrm?resourcekey=0-ajl3Ama1eI_rW1OWtRQbsQ&usp=drive_link

[NValsted]

I can reproduce the problem on Google Colab with the gist from a previous, now-closed, issue about this: #4967 (comment)

At the time of writing, the default environment in Google Colab uses tensorflow=2.18.0 (and keras=3.8.0), but mediapipe-model-maker is constrained to tensorflow<2.16

mediapipe/mediapipe/model_maker/requirements.txt

Line 5 in c39fd6c

tensorflow>=2.10,<2.16

After installing mediapipe-model-maker on the Google Colab instance, either:

• The session must be restarted to use the downgraded versions of tensorflow, keras and others. After this, no GPU can be found
• If the session is not restarted, mediapipe-model-maker throws an exception on import (which I believe is because of the breaking changes introduced with keras 3.0)

[PranayLohia]

@NValsted were you able to fix it?

[NValsted]

I could temporarily resolve the issue by using the fallback runtime in Google colab - see this issue: googlecolab/colabtools#4214. But it's not a proper fix, and ideally mediapipe-model-maker should support the newest versions of the relevant libraries