[FIXED] Tensorflow Module Not Found Error when running MNIST TPU

Issue

I'm new to tensorflow and was trying to follow the tutorial here https://cloud.google.com/tpu/docs/quickstart to run the MNIST TPU model.

I got error from

  python3 mnist_main.py \
  --tpu=$TPU_NAME \
  --model_dir=$MODEL_DIR \
  --data_dir=$DATA_DIR \
  --train_epochs=10 \
  --distribution_strategy=tpu \
  --download

Error message

  Traceback (most recent call last):
  File "mnist_main.py", line 31, in <module>
    from official.vision.image_classification.resnet import common
  File "/usr/share/models/official/vision/image_classification/resnet/common.py", line 26, in <module>
    import tensorflow_model_optimization as tfmot
ModuleNotFoundError: No module named 'tensorflow_model_optimization'

Any idea about to solve it ? Thank you.

Solution

You need to install it using command line (or terminal).

pip install --user --upgrade tensorflow-model-optimization

Answered By - Sean Payne