TensorFlow中如何循环遍历Keras网络数组？

报错根源 #

在tf.function的图执行模式下，tf.range(N)生成的n是Tensor类型的符号化节点，无法直接用来索引Python列表（比如你的networks_array）——Python列表的索引必须是Python整数，而Tensor在图构建阶段无法被解析为具体整数，因此触发TypeError: list indices must be integers or slices, not Tensor。

解决方案 #

根据N是静态已知（定义tf.function时固定）还是动态可变（运行时才确定），分两种处理方式：

情况1：N是静态已知的整数 #

直接用Python原生range循环替代tf.range，TensorFlow的Autograph会自动将静态循环转换为图模式兼容的代码，既保证效率，又能正常访问列表中的模型：

@tf.function
def aggregate_network_outputs(input_data, networks_array, N, K):
    networks_outputs = []
    # 用Python原生range循环，而非tf.range
    for n in range(N):
        input_data_n = tf.concat([input_data[:, n], input_data[:, N:N+K]], axis=1)
        network_output = networks_array[n].some_tf_function(input_data_n)
        networks_outputs.append(network_output)
    return tf.stack(networks_outputs)

情况2：N是动态可变的Tensor #

使用tf.while_loop配合tf.switch_case，通过分支选择的方式在图模式下访问列表中的模型，兼容动态N的场景：

# 提前生成各模型对应的处理分支函数
def build_network_branches(networks_array):
    branches = []
    for idx, _ in enumerate(networks_array):
        def branch(input_data, idx=idx):
            return networks_array[idx].some_tf_function(input_data)
        branches.append(branch)
    return branches

@tf.function
def aggregate_network_outputs(input_data, networks_array, N, K):
    branches = build_network_branches(networks_array)
    # 根据实际输出形状调整element_shape
    networks_outputs = tf.TensorArray(
        size=N, 
        dtype=tf.float32, 
        element_shape=tf.TensorShape([input_data.shape[0], ...])  # 替换...为输出的特征维度
    )
    current_idx = tf.constant(0, dtype=tf.int32)

    def loop_body(idx, outputs):
        # 构造当前模型的输入
        input_data_n = tf.concat([input_data[:, idx], input_data[:, N:N+K]], axis=1)
        # 根据索引选择对应模型的处理逻辑
        current_output = tf.switch_case(idx, branches, default=lambda: tf.zeros_like(input_data_n))
        outputs = outputs.write(idx, current_output)
        return idx + 1, outputs

    # 执行循环
    _, final_outputs = tf.while_loop(
        cond=lambda idx, *args: idx < N,
        body=loop_body,
        loop_vars=[current_idx, networks_outputs],
        shape_invariants=[current_idx.get_shape(), tf.TensorArrayShape(dtype=tf.float32, size=N)]
    )

    return final_outputs.stack()

注意事项 #

• 确保some_tf_function已经被@tf.function装饰，保证其在图模式下能正常运行。
• 如果模型数量极大，tf.switch_case的分支开销可能会显现，此时可以考虑将所有模型的参数合并为大Tensor，通过tf.gather选择参数进行计算，但这种实现复杂度较高，仅适合极端场景。

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