Severe SavedModel load slowdown in tensorflow:24.05+ due to protobuf falling back to Python implementation

AI & Data Science NGC Catalog Feedback
1

Hi NGC TensorFlow team,

In nvcr.io/nvidia/tensorflow:24.05-tf2-py3 and later, we observe a severe slowdown when loading large TensorFlow SavedModels.

Investigation shows that protobuf falls back to the pure-Python implementation:

from google.protobuf.internal import api_implementation
print(api_implementation.Type())  # "python"

This results in roughly a 3× slowdown in SavedModel load time compared to 24.04 and TensorFlow official images where the upb (C++) backend is used.

As a workaround, reinstalling protobuf enables the upb (C++) backend and restores normal performance:

pip install --force-reinstall protobuf==4.25.8

After this, api_implementation.Type() returns "upb".

This suggests a protobuf packaging / build issue in the 24.05+ NGC image.
A fix at the image level (or documentation of this change) would be greatly appreciated.

Thanks for your work on the NGC containers.