如何計算ndarray上一維的平方和？

Question

一個具有形狀的二維矩陣（2,50,25,3）：

xx = np.random.randn(2, 50, 25, 3)

我想計算最后一個維度的平方和。 結果應該是具有形狀（2,50,25,1）的矩陣。

[np.sum(x) for x in np.square(features_displacement[0][0][:])[:]]

此代碼可以成功計算一維，輸出一個形狀為（25,1）的列表，但如何計算如上所述的所有維度？

Answer 1

您可以沿所需的軸應用numpy函數，例如：

np.sum(np.square(xx), axis=3)

將產生一個形狀的陣列(2, 50, 25) 。 不完全確定這是你想要的，如果沒有請更具體:-)

Answer 2

用這個：

sum_of_squares = np.sum(np.square(features_displacement), axis=-1)