在ndarray中获得跨通道元素总和的优化方法

Question

我有一个4D (68000, 96, 124, 4)张量，我想获取所有通道中元素的总和。下面是我正在使用的代码，它很慢。我想知道是否有优化的方法来完成此任务。

下面是我正在使用的代码。

SumChannels = []
for Element in range(0,TrainData.shape[0]):
    row = []
    for channel in range(0,4):
       row.append(TrainData[:,:,:,channel][Element].sum())
    SumChannels.append(row)

Answer 1

可能很简单

SumChannels = TrainData.sum(axis=0)