TensorFlow自定义层含for循环时数据类型异常问题求助

自定义类代码如下：

class MyCalcSqureLayer(keras.layers.Layer):
    def __init__(self, shape ):
        super(MyCalcSqureLayer, self).__init__()
        self.total = self.add_weight(initializer="zero", shape=shape, trainable=False)

        
    def call(self,inputs):
        print('type=',type(inputs))
        x =inputs.shape[0]
        print(x,type(x))
        for i in range(x): # 存在该行时变量x的类型会改变
            print(i)
                
        return inputs

运行上述代码时，输出结果如下：

type= <class 'tensorflow.python.framework.ops.SymbolicTensor'>
None <class 'NoneType'>

但注释掉for i in range(x):该行后，运行结果为：

type= <class 'tensorflow.python.framework.ops.EagerTensor'>
5 <class 'int'>

环境信息：

• tensorflow 2.17.0
• python 3.11.8
• 设备：MacOS M2

需求：希望获取int类型的x，并编写后续的张量操作代码。

解决方案 #

这是TensorFlow图模式与即时执行模式切换导致的：

• 当代码中出现for i in range(x)这类依赖Python原生控制流的代码时，TensorFlow会自动切换到图模式（SymbolicTensor），此时inputs.shape[0]返回静态未知维度值None。
• 注释掉该行后，代码保持在即时执行（Eager Execution）模式，inputs.shape[0]能直接获取到具体的int值。

要稳定获取int类型的维度值，推荐以下方法：

方法1：使用tf.shape()获取动态维度

tf.shape()在两种模式下都能返回具体维度值，返回的EagerTensor可通过.numpy()转为int：

def call(self,inputs):
    print('type=',type(inputs))
    x = tf.shape(inputs)[0].numpy() # 获取动态维度并转成int
    print(x,type(x))
    for i in range(x):
        print(i)
    return inputs

方法2：用TensorFlow原生操作替代Python循环

如果后续需要基于维度做循环操作，优先使用TensorFlow原生API避免触发图模式维度异常：

def call(self,inputs):
    print('type=',type(inputs))
    x = tf.shape(inputs)[0]
    print(x,type(x))
    # 用TensorFlow原生循环替代Python循环，适配图模式
    tf.while_loop(lambda i: i < x, lambda i: (i+1,), (0,))
    return inputs

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