広告

NumPyのarray.sizeに相当するのはPytorchのTensor.numel()


タグ:python machine_learning

テンソルの次元はベクトル(Python的に言えばタプル)で表されるが、 それとは別に要素数(結局スカラが何個入っているのか?)が欲しいことがある。
つまり例えばC×H×Wのテンソルに対しタプル (C, H, W) ではなく掛け算した後の CHWなる値が直接欲しい場合だ。

これ自体はNumPyでもPyTorchでもできるが、微妙にメソッド名が違いややこしかった。
結論としては
  • NumPy: np.array.size
  • PyTorch: torch.Tensor.numel()
でできる。

使用例:
  In [1]: import numpy as np
  In [2]: from torch import Tensor

  In [3]: np.asarray([[1, 2, 3]]).shape  # これは形状
  Out[3]: (1, 3)

  In [4]: np.asarray([[1, 2, 3]]).size  # こっちは要素数
  Out[4]: 3

  In [5]: Tensor([[1, 2, 3]]).size  # 関数なのでカッコで呼び出すこと
  Out[5]: 

  In [6]: Tensor([[1, 2, 3]]).size()  # 独自型だしshape相当ではないか
  Out[6]: torch.Size([1, 3])

  In [7]: Tensor([[1, 2, 3]]).numel()  # これが欲しかったのだ
  Out[7]: 3

ちなみに"numel"は"number of elements (要素数)"の略と思われる。なぜそろえなかったのか。