PyTorch で tensor をある軸に従って flip する方法

目的

seq2seqなどのrnnモデルで学習する時、答えのシーケンスを逆転する必要があります。普通のnumpyでやればa[::-1]でできるが、pytorch がマイナスインデクシングをまだ実装してないのでやり方をまとめます。

numpyの場合

>>> import numpy as np
>>> a=np.random.rand(2,3,4)
>>> a
array([[[0.77035401, 0.24779618, 0.76642706, 0.05111546],
        [0.93068699, 0.63876274, 0.37042163, 0.09062927],
        [0.91439172, 0.91490527, 0.82236844, 0.74023026]],

       [[0.0294182 , 0.15135825, 0.22691598, 0.64636446],
        [0.58836392, 0.97549164, 0.39324753, 0.66243478],
        [0.38082428, 0.20068817, 0.42776323, 0.49329102]]])
>>> a[:,:,::-1]
array([[[0.05111546, 0.76642706, 0.24779618, 0.77035401],
        [0.09062927, 0.37042163, 0.63876274, 0.93068699],
        [0.74023026, 0.82236844, 0.91490527, 0.91439172]],

       [[0.64636446, 0.22691598, 0.15135825, 0.0294182 ],
        [0.66243478, 0.39324753, 0.97549164, 0.58836392],
        [0.49329102, 0.42776323, 0.20068817, 0.38082428]]])

pytorch の場合

torch.arange(3,-1,-1) を使うことで、インデックスの逆転を通じてテンサーの軸を逆転します。

>>> a = torch.rand(2,3,4)
>>> a
tensor([[[0.7077, 0.0939, 0.3385, 0.8469],
         [0.1563, 0.1077, 0.5764, 0.2502],
         [0.3548, 0.1774, 0.8615, 0.1192]],

        [[0.0571, 0.6192, 0.9620, 0.3761],
         [0.3884, 0.6482, 0.3567, 0.4959],
         [0.6242, 0.9567, 0.0100, 0.9072]]])
>>> a[:,:,torch.arange(3,-1,-1)]
tensor([[[0.8469, 0.3385, 0.0939, 0.7077],
         [0.2502, 0.5764, 0.1077, 0.1563],
         [0.1192, 0.8615, 0.1774, 0.3548]],

        [[0.3761, 0.9620, 0.6192, 0.0571],
         [0.4959, 0.3567, 0.6482, 0.3884],
         [0.9072, 0.0100, 0.9567, 0.6242]]])