TensorFlow并行通道模型连接异常及子模型构建报错排查

• 调试支持可变数量RGB图像输入通道的自研TensorFlow/Keras模型时，疑似存在通道未正确连接问题
• 实例化2通道模型调用m.summary()查看结构，发现2个SlicingOpLambda切片算子均连接到input_25[0][0]，疑似两个切片取了相同通道数据，未按索引切分不同通道
• 构建子模型验证单通道分支输入输出时，执行r_m = tf.keras.Model(model.inputs, model.layers[3].input)可正常运行，执行r_m2 = tf.keras.Model(model.inputs, model.layers[3].output)时抛出ValueError，提示Graph disconnected（计算图断开），无法获取rescaling_4层对应的输入张量

IMG_SHAPE = (160, 160, 3)

def get_ch_model_simple():
    i_input = tf.keras.Input(shape=IMG_SHAPE)    
    # scale pixels to float
    x = tf.keras.layers.Rescaling(1.0 / 255)(i_input)
    x = tf.keras.layers.Conv2D(32, kernel_size=(3,3), activation="relu")(x)
    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)

    return tf.keras.Model(i_input, x)
    
def get_model(n_chan=2):
    
    inputs = tf.keras.Input(shape=(n_chan, 160, 160, 3))

    ch_features = []
    for ch in range(n_chan):
        ch_model = get_ch_model_simple()
        # select specific channel
        ch_model_input = inputs[:,ch,:,:,:]
        i_ch_features = tf.keras.layers.Flatten()(ch_model(ch_model_input))
        i_ch_features = tf.keras.layers.Dropout(0.5)(i_ch_features)
        ch_features.append(i_ch_features)

    all_ch_features = tf.keras.layers.concatenate(ch_features)
    outputs = tf.keras.layers.Dense(2, activation = "softmax")(all_ch_features)
    
    return tf.keras.Model(inputs, outputs)

• 切片操作的循环变量捕获错误：Python for循环中的ch变量为延迟引用，TensorFlow构建静态计算图时不会在单次循环迭代时锁定当前ch的数值，所有切片算子最终都会引用循环结束后ch的最终值，导致两个切片实际读取同一索引的通道数据，对应summary中所有切片算子连接到同一输入节点的现象。
• 计算图断开源于子模型拓扑不连通：循环中每次调用get_ch_model_simple()都会生成独立的Functional子模型，子模型自带独立的Input占位层。直接将切片张量传入ch_model()时，张量会直接喂给子模型内部的Rescaling层，跳过了子模型自身的Input层。当尝试直接取Rescaling层输出构建新模型时，Keras沿计算图反向追溯会发现Rescaling层的上游是未接入外层模型的子模型Input层，无法形成从外层输入到目标输出的完整路径，因此抛出Graph disconnected错误。

1. 修复切片变量捕获问题：使用tf.keras.layers.Lambda层封装切片操作，通过默认参数显式绑定当前循环的通道索引，避免静态图构建时的变量引用偏差。
2. 修复计算图连通性问题：确保子模型的Input层与外层计算图完整连通，传入子模型的张量直接对接子模型Input，不跳过子模型的输入节点。如果需要通道间特征提取权重共享，可将单通道子模型的实例化操作移到循环外；如果需要各通道权重独立，保留子模型在循环内实例化即可。

修复后的可运行代码如下：

import tensorflow as tf
IMG_SHAPE = (160, 160, 3)

def get_ch_model_simple():
    i_input = tf.keras.Input(shape=IMG_SHAPE)    
    # 像素值归一化
    x = tf.keras.layers.Rescaling(1.0 / 255)(i_input)
    x = tf.keras.layers.Conv2D(32, kernel_size=(3,3), activation="relu")(x)
    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
    return tf.keras.Model(i_input, x)
    
def get_model(n_chan=2):
    inputs = tf.keras.Input(shape=(n_chan, 160, 160, 3))
    ch_features = []
    # 权重共享则将子模型实例化放在循环外，权重独立则移到循环内
    ch_model = get_ch_model_simple()
    for ch_idx in range(n_chan):
        # Lambda层显式绑定当前通道索引，固定切片参数
        ch_input = tf.keras.layers.Lambda(
            lambda x, c=ch_idx: x[:, c, :, :, :],
            name=f"slice_ch_{ch_idx}"
        )(inputs)
        # 子模型输入直接接切片层输出，完整连通子模型Input节点
        ch_feat = ch_model(ch_input)
        ch_feat = tf.keras.layers.Flatten()(ch_feat)
        ch_feat = tf.keras.layers.Dropout(0.5)(ch_feat)
        ch_features.append(ch_feat)

    all_ch_features = tf.keras.layers.concatenate(ch_features)
    outputs = tf.keras.layers.Dense(2, activation="softmax")(all_ch_features)
    
    return tf.keras.Model(inputs, outputs)

验证说明 #

• 修复后调用m.summary()可看到每个slice_ch_x层对应独立的通道索引，不会再出现所有切片算子连接到同一输入节点的问题
• 构建子模型时，沿拓扑路径从外层输入追溯到目标层输出不会再触发计算图断开错误，例如执行r_m2 = tf.keras.Model(m.inputs, m.get_layer(index=2).get_layer("rescaling").output)可正常运行。

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