【PyTorch】テンソルがraw majorであることの確認

はじめに

　↓の過去記事で「PyTorchでは、テンソルが作成されたとき、その要素はメモリ上にraw majorと呼ばれる方式で並びます」と書きました。本記事では、テンソルの要素がraw majorで並んでいることを、各要素のメモリアドレスを見ることで確認します。

aisinkakura-datascientist.hatenablog.com

raw majorとは？

　raw majorとは行列を図１の上図に記載した順番でメモリに乗せることをいいます。また、raw majorと対応する言葉として、column majorがあります。column majorは図１の下図に記載した順番でメモリに乗せることをいいます。

図１：raw majorとcolumn major

PyTorchのテンソルがraw majorで並んでいることの確認

　テンソルが格納されているメモリの先頭アドレスは、data_ptr()メソッドを使うことで取得できます。なので、下記のコードで全ての要素の先頭メモリアドレスを表示させることができます。

import itertools
import torch

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
for i, j in itertools.product(range(x.size(0)), range(x.size(1))):
  print(f"x[{i}, {j}] address: {x[i, j].data_ptr()}")

　私の環境では、結果は図２のようになりました。

図２：テンソルの要素の先頭メモリアドレス

　したがって、今回定義したテンソルxの各要素は図３のようにメモリ上に乗っています。

図３：テンソルの要素のメモリ上での並び方

　以上より、PyTorchのテンソルがraw majorで並んでいることを確認できました。