Оптимизированный способ получения суммы элементов по каналам в ndarray

Question

У меня есть тензор 4D (68000, 96, 124, 4), и я хочу получить сумму элементов по всем каналам. Ниже приведен код, который я использую, и он медленный. Мне интересно, есть ли оптимизированный способ сделать это.

Ниже приведен код, который я использую.

Пример:

TrainData = np.random.rand(20,5,3,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)

Результатом является SumChannels,вложенный list длиной 20 (из примера) с каждым inner list длиной 4

PS: Ну, я понятия не имею, почему приведенный выше код был слишком медленным, но когда я перезагружал свою машину. Выход был мгновенный