如何优化 NumPy/TensorFlow 中张量运算的嵌套循环？

Question

我正在研究一个机器学习问题，涉及分类任务的蒙特卡罗模拟。我当前的实现涉及基于多项分布生成合成类标签，并计算输入特征和生成标签之间的特定张量积。目标是估计 ANN 分类模型的正则化阈值。然而，用于计算四维数组 xy 的嵌套循环似乎是一个瓶颈，我正在寻找使用矢量化或任何其他有效方法来优化它的方法。

当前实施：

import numpy as np
import tensorflow as tf

def lambda_qut_sann_classification(xsample, hat_p_training, nSample=100000, miniBatchSize=500, alpha=0.05, option='quantile'):
    if np.mod(nSample,miniBatchSize) == 0:
        offset = 0
    else: 
        offset = 1
        
    n, p1 = xsample.shape
    number_class = len(hat_p_training)

    fullList = np.zeros((miniBatchSize*(nSample//miniBatchSize+offset),))

    for index in range(nSample//miniBatchSize+offset):
        ySample = np.random.multinomial(1, hat_p_training, size=(n, 1, miniBatchSize))
        y_mean = np.mean(ySample, axis=0)

        xy = np.zeros(shape=(n, p1, miniBatchSize, number_class))
        for index_n in np.arange(n):
            xy[index_n, :, :, :] = np.outer(xsample[index_n, :], (y_mean-ySample)[index_n, :, :]).reshape((p1, miniBatchSize, number_class))
        
        xymax = np.amax(tf.reduce_sum(np.abs(tf.reduce_sum(xy, axis=0).numpy()), axis=2).numpy(), axis=0)

    # Further processing...

具体问题：

• 如何优化 xy 的计算，目前使用 for 循环和整形操作实现，可能使用 NumPy 或 TensorFlow 中的矢量化技术？目标是消除或减少 for 循环以提高效率。
• 是否有更有效的方法来执行这些张量运算，可以利用 TensorFlow 或 NumPy 的功能来获得更好的性能？

背景：

• xsample 是形状 (n, p1) 的 2D NumPy 数组，表示输入特征。
• hat_p_training 是一个一维数组，表示估计的类别概率。
• 代码生成 ySample（基于 hat_p_training 的合成标签集），并计算 xy（表示转置 xsample 与差值 y_mean-ySample 之间的乘积的张量）。
• 最终目标是找到 xy 某些维度上的最大值以进行进一步处理。

我很感激任何关于优化这部分代码的见解或建议。谢谢！

Answer 1

使用

outer

的内部循环可以替换为单个广播乘法。

进行示例计算：

In [42]: n, p1, miniBatchSize, number_class = 5,4,3,2
    ...: x = np.arange(n*p1).reshape(n,p1)
    ...: y  = np.arange(n*miniBatchSize*number_class).reshape(n,miniBatchSize,number_class)
    ...: xy = np.zeros(shape=(n, p1, miniBatchSize, number_class))
    ...: for index_n in np.arange(n):
    ...:     xy[index_n, :, :, :] = np.outer(x[index_n, :], y[index_n, :, :]).reshape((p1, miniBatchSize, number_class))
    ...:     

In [43]: xy.shape
Out[43]: (5, 4, 3, 2)

以及广播的等效内容：

In [44]: xy1 = x[:,:,None,None] * y[:,None,:,:]

In [45]: xy1.shape
Out[45]: (5, 4, 3, 2)
In [46]: np.allclose(xy, xy1)
Out[46]: True

取决于

n

的大小，这可以节省时间，并且也可以更容易“矢量化”外循环。我没有尝试理解或重写那部分。