Cannot set BatchNormalization's datatype to `float64` by setting _dtype

[maybeLee]

Please go to TF Forum for help and support:

https://discuss.tensorflow.org/tag/keras

If you open a GitHub issue, here is our policy:

It must be a bug, a feature request, or a significant problem with the documentation (for small docs fixes please send a PR instead). The form below must be filled out.

Here's why we have that policy:.

Keras developers respond to issues. We want to focus on work that benefits the whole community, e.g., fixing bugs and adding features. Support only helps individuals. GitHub also notifies thousands of people when issues are filed. We want them to see you communicating an interesting problem, rather than being redirected to Stack Overflow.

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.8.2
• Python version: 3.7
• Bazel version (if compiling from source): N/A
• GPU model and memory: N/A
• Exact command to reproduce: Please see the code snippet or colab link below.

You can collect some of this information using our environment capture script:

https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh

You can obtain the TensorFlow version with: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the problem. I want to automatically changing the datatype for some layers inside a DL model. By searching the internet, I was told that changing _dtype property of a layer is feasible to do so. That's why I want to change the _dtype of some layers. However, I can change almost all layers' datatype except for the BatchNormalization layer. Specifically, I find that I cannot change the datatype of BatchNormalization layer to float64 or double but I can change it to bfloat16, float16, float32.

Describe the current behavior. Currently, I cannot change the datatype of BatchNormalization to float64 or double by setting the _dtype property of BatchNormalization, but I can change the datatype to bfloat16, float16, float32. I can also change the datatype of other layers such as Conv2D.

Describe the expected behavior. It would be of great help if I can also change the datatype of BN to float64 by setting _dtype. Currently, this strategy seems buggy.

Contributing.

• Do you want to contribute a PR? (yes/no): yes
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue.

import keras
import traceback

def try_dtype_bn(dtype):
  import numpy as np
  test_input = np.random.rand(10,10,10,3)

  x = keras.layers.Input((10,10,3), dtype=dtype)
  try:
    layer = keras.layers.BatchNormalization(dtype=dtype)
    y = layer(x)
    model = keras.models.Model(x,y)
    res = model.predict(test_input)
    print(res.shape, res.dtype)
    print(f"Success when directly setting BN's datatype to {dtype} through function call")
  except:
    print(f"Fail when directly setting BN's datatype to {dtype} through function call")
    print(traceback.format_exc())

  try:
    layer = keras.layers.BatchNormalization()
    layer._dtype=dtype
    y = layer(x)
    model = keras.models.Model(x,y)
    import numpy as np
    test_input = np.random.rand(10,10,10,3)
    res = model.predict(test_input)
    print(res.shape, res.dtype)
    print(f"Success when directly setting BN's datatype to {dtype} by setting _dtype property")
  except:
    print(f"Fail when setting BN's datatype to {dtype} by setting _dtype property")
    print(traceback.format_exc())

try_dtype_bn("float64")
try_dtype_bn("half")

You can also access the colab notebook here: https://colab.research.google.com/drive/1FZ8_DNWDYDeSa11VLAMAUK4xx1qp2k7n?usp=sharing

Source code / logs.

Traceback (most recent call last):
  File "<ipython-input-28-7b93c8890c00>", line 23, in try_dtype_bn
    y = layer(x)
  File "/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/op_def_library.py", line 57, in _SatisfiesTypeConstraint
    f"Value passed to parameter '{param_name}' has DataType "
TypeError: Exception encountered when calling layer "batch_normalization_36" (type BatchNormalization).

Value passed to parameter 'x' has DataType float64 not in list of allowed values: float16, bfloat16, float32

[sushreebarsa]

@gadagashwini i was able to replicate the issue on colab, please find the gist here. Thank you!

[haifeng-jin]

Setting the private attribute _dtype of a layer is not an official supported way to dynamically change dtypes of layers. To discuss: Is there any official way to dynamically change the dtype of layers?