PyTorch中的最大功能与NumPy中的最大功能

问题描述

我正在开发一种算法，其中涉及允许NumPy的CPU情况，以及允许PyTorch的GPU情况。该对象几乎总是4D。该对象的两个版本如下。

B = [
    [[[0.5000,0.5625],[0.5000,0.5625]],[[1.2500,0.5000],[0.5625,0.6875]],[[0.5625,0.6250],0.5625]]]
] 

B_array = np.array(B)

B_tensor = torch.Tensor(B)

我想获取每个2D矩阵的最大值，以便得到以下结果：

max_array_fn(B_array)   # returns array([0.5625,1.250,0.6250])
max_tensor_fn(B_tensor) # returns tensor([0.5625,0.6250])

这里讨论了部分解决方案，但这仅适用于CPU上的NumPy：

Max of each 2D matrix in 4D NumPy array

但是，在GPU上，PyTorch似乎没有使用与NumPy相同的约定。

如果将其定义为NumPy数组，我们可以使用np.array(B,axis=(0,2,3))解决问题。不支持在PyTorch中使用类似于NumPy示例的axis，如下所示：

PyTorch torch.max over multiple dimensions

该方法是否有替代向量化？为什么只能使用NumPy在CPU上对其进行矢量化，而不能使用PyTorch在GPU上进行矢量化？

正确的解决方案不应使用任何循环，最好不要使用一个函数调用。

解决方法

不是最优雅的方式：

B.max(dim=3)[0].max(dim=2)[0].max(dim=0)[0]