pytorch 统计指定维度上的元素数目 torch.size numel

需求；对给定的tensor, 获取其指定几个维度上的张量元素数目。

如：对于tensor: $x\in {\mathbb{R}}^{d_0\times d_1\times d_2\times d_3}$，返回$d_0\times d_1$.

方法：

torch.size(x)[:2].numel()

其中torch.tensor.numel()为统计元素数目的函数，当输入为torch.tensor类型时，返回该tensor的元素数目；当输入为torch.size类型时，返回torch.size的各项乘积。如：

a = torch.randn(2,3,4,5)
torch.numel(a)
120 # (= 2*3*4*5)

torch.size(a).numel()
120 # (= 2*3*4*5)