【Pytorch】torch.sumのdim=-1, 0, (1, 1)などの意味とは？実際にコードを動かして検証してみた

githubでコードを見ていると、torch.sum(m, dim=-1)のように傍目では意味が分からない部分が出てきたので、torch.sum()について色々と調べてみました。

torch.sumの役割

torch.sum()の役割は、与えられたテンソルの要素の合計を返すことです。
これは、sumという単語から想像できますね。

torch.sumの引数

引数は、torch.sum(input, dim, keepdim=False, *, dtype=None) です。
それぞれの意味を見ていきます。

dimとkeepdimの説明が良く分からないと思うので、実際にコードを動かして調べてみました。（私も説明文だけでは分かりませんでした）

dimの詳しい説明

dimを設定しない場合

まずは、dimを設定しなかった場合にどのような働きをするのか見ていきます。

m = torch.arange(16).reshape(4, -1)
"""
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
"""

torch.sum(m)
# tensor(120)

dimを設定しない場合は、テンソルのすべての要素の合計が返ってくることが分かります。

dimをint型で設定する場合

m = torch.arange(16).reshape(4, -1)
"""
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
"""

次に、dimをint型で設定する場合にどのような働きをするのか見ていきます。
上記のテンソルの場合に設定できるdimは、-2~1の整数です。
それ以外の整数を入れると、次のようなエラーが出ます。
IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)

上記のエラーが出る理由は、mの次元数が２でそれ以上の値を設定したからです。
入力テンソルの次元数を増やせば、入れられる整数の範囲も広がります。式で表すと、-(入力テンソルの次元数)~(入力テンソルの次元数-1)になります。
今回は、分かりやすくするために4×4の２次元のテンソルの場合を見ていきます。

# 以下のテンソルをもとにそれぞれのdimの場合を見ていきます
m = torch.arange(16).reshape(4, -1)
"""
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
"""

torch.sum(m, dim=1)
# tensor([ 6, 22, 38, 54])

torch.sum(m, dim=0)
# tensor([24, 28, 32, 36])

torch.sum(m, dim=-1)
# tensor([ 6, 22, 38, 54])

torch.sum(m, dim=-2)
# tensor([24, 28, 32, 36])

dimをtuple型で設定する場合

続いて、dimをtuple型で設定する場合を見ていきます。
同じように以下のテンソルの場合で見ていきます。

m = torch.arange(16).reshape(4, -1)
"""
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
"""

torch.sum(m, dim=(1, ))
# tensor([ 6, 22, 38, 54])

torch.sum(m, dim=(1, 0))
# tensor(120)

keepdimの詳しい説明

keepdimを設定した場合にどのような挙動をするのか見ていきます。

torch.sum(m, dim=1, keepdim=True)
"""
tensor([[ 6],
        [22],
        [38],
        [54]])
"""
# torch.Size([4, 1])

まとめ

あわせて読みたい

【Pytorch】torch.catの意味や引数を詳しく解説！dim=-1を設定した場合は？ torch.catはgithubでモデルの構成を見ていると、かなり登場するので色々と調べてみました。例を使って色々と試しています。 【の意味】 torch.catの役割は、与えられた...

あわせて読みたい

【PyTorch】ベクトルを操作する代表的な関数を解説！view, reshape, transpose, unsqueeze, matmul, ein... PyTorchを使っているコードを見ているとベクトルを操作する関数がたくさん出てくると思います。しかし、どんな機能なのかわからないと自分でコードを変更しようとしても難しいです。そこで今回は、実際にベクトルを操作しながらそれぞれの関数を解説していきます。

あわせて読みたい

【Pytorch】nn.Linearの引数・ソースコードを徹底解説！ torch.nn.Linearは基本的な全結合層です。今回は、nn.Linearの引数とソースコードをしっかりと説明していきます。曖昧な理解を直していきましょう。