Change dtype policy of tensorflow saved model

user_5656827 0
Viewed 471

I have trained a model using tensorflow 2.2.0 with a mixed_float16 policy. Training and evaluation of the model went fine, but now the model cannot be evaluated on devices that do not support mixed_float16, e.g. older GPUs or CPUs.

For example, when loading the model on a computer without GPU I get the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: No OpKernel was registered to support Op 'MaxPool3D' used by {{node model/max_pooling3d/MaxPool3D}} with these attrs: [ksize=[1, 2, 2, 2, 1], padding="SAME", T=DT_HALF, data_format="NDHWC", strides=[1, 2, 2, 2, 1]]
Registered devices: [CPU]
Registered kernels:
  device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_BFLOAT16, DT_HALF]
  device='XLA_GPU_JIT'; T in [DT_FLOAT, DT_BFLOAT16, DT_HALF]
  device='CPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_HALF]
  device='GPU'; T in [DT_FLOAT]

     [[model/max_pooling3d/MaxPool3D]] [Op:__inference_predict_function_19293]

I have tried to set the TensorFlow policy locally before loading the model via

policy = tensorflow.keras.mixed_precision.experimental.Policy("float32")
tensorflow.keras.mixed_precision.experimental.set_policy(policy)

Unfortunately, that does not change the dtype of the model.

Is there anyway to change the dtype of a changed model? Or to somehow save the model to allow it to automatically choose for float32 when float16 is not available?

I've also tried to look into using TFLite for this purpose, but the optimization there can also not be done due to issues with the model being saved as in mixed precision format.

tensorflow python
1 Answers

I had tried to save the weights and for model and then load them into a newly constructed model, but that didn't work. However, directly copying the weights from the trained model to a new model directly. In short, something like this:

import tensorflow as tf
trained_model = tf.keras.models.load_model(/model/path)

# Function to create the model architecture without trained weights
# Make sure that this model is not created using a mixed precision policy!
new_model = create_model()

new_model.set_weight(trained_mode.get_weights())

The new_model can then be used on CPU as well.

user_56568270
Related