Keras中矩阵按列操作的正确实现方法咨询

这个问题我之前踩过坑！你遇到的核心问题是Keras的符号张量和Python循环变量不兼容——在构建计算图的时候，Python循环里的i是普通的Python变量，等循环结束后它会固定为最后一个迭代值，所以你的Lambda层其实自始至终都在取矩阵的最后一列，运行时自然每一步拿到的结果都一样。

下面给你几种Keras/TensorFlow里按列操作的正确姿势：

方法1：用tf.map_fn（最直观，通用场景） #

tf.map_fn专门用来对张量的某个维度逐一应用自定义函数，完美适配按列处理的需求。我们可以先把矩阵转置，让列变成行（因为map_fn默认处理第0维度），然后遍历每一行（原列）和v计算函数f：

import tensorflow as tf

# 先定义你的函数f，比如这里以计算v和列的点积为例
def f(v, column):
    return tf.reduce_sum(v * column, axis=0)

# 包装一个闭包，让内部函数能捕获v这个张量
def process_column(v):
    def inner(col):
        return f(v, col)
    return inner

# 假设v是形状为(D,)的张量，M是形状为(D, N)的矩阵
transposed_M = tf.transpose(M, perm=[1, 0])  # 转置后形状变为(N, D)，每一行对应原矩阵的一列
# 对转置后的每一行应用函数，得到每个列的计算结果
column_results = tf.map_fn(process_column(v), transposed_M)
# 如果需要结果是列向量形式，可以再调整形状
final_result = tf.expand_dims(column_results, axis=0)

方法2：用tf.split拆分列后批量处理 #

先把矩阵按列拆分成单独的张量列表，再逐个和v计算f，最后把结果拼接起来：

import tensorflow as tf

# 动态获取M的列数（比静态的K.int_shape更适配动态形状场景）
num_columns = tf.shape(M)[1]
# 按axis=1拆分，得到N个形状为(D, 1)的张量
column_list = tf.split(M, num_or_size_splits=num_columns, axis=1)

# 对每个列张量应用f，注意可以用tf.squeeze去掉多余的维度
results = [f(v, tf.squeeze(col, axis=1)) for col in column_list]
# 把所有结果拼接成形状为(N,)的张量（或按需调整为矩阵）
final_result = tf.stack(results, axis=0)

方法3：向量化运算（性能最优，优先考虑） #

如果你的函数f可以通过TensorFlow的广播机制实现向量化计算，那这是最快的方式，完全不需要循环：
比如f是v和列的点积，直接用矩阵乘法就能搞定：

import tensorflow as tf

# v形状(D,)，M形状(D,N)，点积结果直接是v与M的矩阵乘法
final_result = tf.matmul(tf.expand_dims(v, axis=0), M)
# 或者用更简洁的tensordot
final_result = tf.tensordot(v, M, axes=1)

这种方法没有循环开销，性能最好，能向量化就优先用它！

总结一下：别再用Python循环+Lambda取列了，Keras的符号张量需要用TensorFlow提供的符号化操作来处理维度遍历，这样才能保证计算图构建正确，运行时拿到预期结果。

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