Note: the content in this post was tested with the GPU version of TensorFlow 2.0.

I need float64 for my models. However, I did not realize until recently that even though the dtype of all layers in my models is tf.float64, and even though my x is a tf.float64 tensor, when I call model(x), the model still treats x as a tf.float32 tensor.

Here’s a simple example to check this:

1
class TestModel(tf.keras.Model):
2
    def __init__(self):
3
        super(TestModel, self).__init__()
4
    def call(self, inputs, training=False):
5
        print("The type of input: {}".format(inputs.dtype))
6
        return inputs

And then, in any Python interpreter or script:

1
x = tf.constant(1.0, dtype=tf.float64)
2
print("The type of x: {}".format(x.dtype))
3
model = TestModel()
4
y = model(x)
5
print("The type of output: {}".format(y.dtype))

In the output, we should see something like:

1
The type of x: <dtype: 'float64'>
2
The type of input: <dtype: 'float32'>
3
The type of output: <dtype: 'float32'>

To solve this issue, we also have to provide the dtype argument to our model. Something like:

1
class TestModel(tf.keras.Model):
2
    def __init__(self, dtype=tf.float64):
3
        super(TestModel, self).__init__(dtype=dtype)
4
    def call(self, inputs, training=False):
5
        print("The type of input: {}".format(inputs.dtype))
6
        return inputs

Now we can see that x, inputs, and y are all tf.float64.

1
The type of x: <dtype: 'float64'>
2
The type of input: <dtype: 'float64'>
3
The type of output: <dtype: 'float64'>

This is not a complicated problem. But I was not aware of it because the official documentation did not mention it, at least not at the time I wrote this post. And I did not get any warning or error message during runtime either.

The underlying reason is that in a parent class of tf.keras.Model, namely tf.keras.layers.Layer, the __call__ method calls a casting method, _maybe_cast_input. And _maybe_cast_input will cast inputs to the model’s dtype.

So we must specify the model’s dtype as the one we actually want. Otherwise, no matter how we set up our data or layers, the model will always use tf.float32 by default.

Another workaround is to force TensorFlow to use tf.float64 as the default floating-point type. We can do that through tf.keras.backend.set_floatx('float64').