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

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

如：对于tensor: x ∈ R d 0 × d 1 × d 2 × d 3 xin mathbb R ^{d_0times d_1times d_2 times d_3 } x∈Rd0​×d1​×d2​×d3​，返回 d 0 × d 1 d_0 times d_1 d0​×d1​.

方法：

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)