如何用NumPy高效计算张量Tp,q,r = Sp,q·Sp,r？

你的需求本质是对每个序列计算外积（即序列中每个元素与自身所有元素相乘，生成二维矩阵），完全可以通过Numpy的广播机制去掉所有循环，一步到位完成计算，效率比原代码高得多。

优化后的代码 #

import numpy as np

# 将输入的序列列表转为Numpy数组，shape为(nseqs, seq_length)
seq_array = np.array(seq)
# 利用广播直接计算所有序列的外积，得到shape为(nseqs, seq_length, seq_length)的结果
output = seq_array[:, :, np.newaxis] * seq_array[:, np.newaxis, :]

原理说明 #

• seq_array[:, :, np.newaxis]：给原数组增加一个维度，shape变为(nseqs, seq_length, 1)，相当于把每个序列的元素都变成单列向量。
• seq_array[:, np.newaxis, :]：给原数组在中间增加一个维度，shape变为(nseqs, 1, seq_length)，相当于把每个序列的元素都变成单行向量。
• Numpy的广播机制会自动把这两个数组扩展到相同的(nseqs, seq_length, seq_length)形状，对应位置相乘就正好得到tₚ,q,r = sₚ,q·sₚ,r的结果。

对比原代码的优势 #

原代码通过两层循环复制元素再转置相乘，本质和外积逻辑一致，但Python循环的效率远低于Numpy底层的C实现优化，尤其是当nseqs或seq_length较大时，优化后的代码速度会有显著提升，同时代码更简洁易读。

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