行列の積

PyTorch を用いたバッチへの行列の積がブロードキャストされた結果ってどうなんだっけ？というのが今更に不安になったので実験。numpy も使う。

#! /usr/bin/env python
# -*- coding: utf-8 -*-

import torch

def main():
    batch = torch.Tensor([
        [
            [0, -1],
            [1,  0]
        ],
        [
            [1, 0],
            [0, 1]
        ],
        [
            [1, 2],
            [3, 4]
        ],
    ])

    W = torch.Tensor([
        [1, -1],
        [-1, 1]
    ])

    print(W @ batch)

if __name__ == "__main__":
    main()

というだけのテスト。import torch を import numpy as np にして numpy 版は試す。

結果

numpy

[[[-1 -1]
  [ 1  1]]

 [[ 1 -1]
  [-1  1]]

 [[-2 -2]
  [ 2  2]]]

PyTorch

tensor([[[-1., -1.],
         [ 1.,  1.]],

        [[ 1., -1.],
         [-1.,  1.]],

        [[-2., -2.],
         [ 2.,  2.]]])

普通に “当たり前” の結果で良かった。O'Reilly Japan - ゼロから作るDeep Learning ❸ でももう一度読まないと backprop で W がどうやっていい感じに update されるのか分からなくて怖すぎ・・・。って言うか、broadcast で W ってどういうことになるの・・・。

＊＊＊＊＊

O'Reilly Japan - ゼロから作るDeep Learning ❸ の p.299「ステップ 40 ブロードキャストを行う関数」がこの疑問に対応する。特に pp.304-305 が関連する。当たり前ではあるが、フレームワークのほうでその辺はよろしく対応してくれているということである。