TensorFlow 2.7.0中tf.cast数据类型未更改的异常问题

问题背景 #

在TensorFlow 2.7.0中使用tf.cast时，因对张量类型判断的误解，导致模型导出到TensorFlow.js后，Conv2D层因输入类型为int32而报错。

运行以下代码：

import tensorflow as tf
import numpy as np

inputs = tf.constant(np.array([[0., 0., -1., 0., 0., 0., 0., -1., 1.]]), dtype=tf.float32)
print(tf.shape(inputs))

x = tf.cast(inputs + 1, tf.int32)
print(tf.shape(x))
x = tf.one_hot(x, 3)
print(tf.shape(x))
x = tf.cast(x, tf.float32)
print(tf.shape(x))

得到输出：

tf.Tensor([1 9], shape=(2,), dtype=int32)
tf.Tensor([1 9], shape=(2,), dtype=int32)
tf.Tensor([1 9 3], shape=(3,), dtype=int32)
tf.Tensor([1 9 3], shape=(3,), dtype=int32)

起初误以为tf.cast未将张量转为float32，但直接打印张量时，可见其dtype确实是float32：

tf.Tensor([1 9], shape=(2,), dtype=int32)
tf.Tensor([[ 0.  0. -1.  0.  0.  0.  0. -1.  1.]], shape=(1, 9), dtype=float32)
tf.Tensor([1 9], shape=(2,), dtype=int32)
tf.Tensor([[1 1 0 1 1 1 1 0 2]], shape=(1, 9), dtype=int32)
tf.Tensor([1 9 3], shape=(3,), dtype=int32)
tf.Tensor(
[[[0. 1. 0.]
  [0. 1. 0.]
  [1. 0. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [1. 0. 0.]
  [0. 0. 1.]]], shape=(1, 9, 3), dtype=float32)
tf.Tensor([1 9 3], shape=(3,), dtype=int32)
tf.Tensor(
[[[0. 1. 0.]
  [0. 1. 0.]
  [1. 0. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [0. 1. 0.]
  [1. 0. 0.]
  [0. 0. 1.]]], shape=(1, 9, 3), dtype=float32)

原因分析 #

• tf.shape()返回的是int32类型的形状张量，仅用于表示输入张量的维度信息，它的dtype与原张量的dtype无关。无论原张量是float32还是int32，tf.shape()的输出始终为int32，这是TensorFlow的设计逻辑，并非tf.cast失效。
• 模型导出到TF.js后Conv2D报错的核心原因，是模型中某个节点的dtype未被正确转换为float32，导致TF.js的Conv2D层（仅支持float32输入）接收到int32类型数据。

解决方法 #

1. 正确判断张量dtype：不要通过tf.shape()的输出判断张量类型，直接使用张量.dtype属性查看，或打印张量本身查看其dtype字段。
2. 显式确保类型转换生效：在模型构建过程中，对需要转换类型的节点，显式保留转换后的张量，避免中间节点残留int32类型：

   x = tf.cast(inputs + 1, tf.int32)
   x = tf.one_hot(x, 3)
   # 显式转换并赋值，确保后续使用的是float32张量
   x = tf.cast(x, tf.float32)
   

3. 导出前验证模型类型：保存模型前，检查模型各层的输入输出dtype，确保Conv2D等层的输入为float32：

   model = tf.keras.Model(inputs=inputs, outputs=x)
   # 查看输入输出dtype
   print("Input dtype:", model.inputs[0].dtype)
   print("Output dtype:", model.outputs[0].dtype)
   

4. TF.js端类型处理：如果导入后仍有类型问题，推理时显式将输入转换为float32：

   const inputTensor = tf.tensor2d([[0., 0., -1., 0., 0., 0., 0., -1., 1.]]).cast('float32');
   const output = await model.predict(inputTensor);
   

内容的提问来源于stack exchange，提问作者Andrea