🔢

PyTorchのbroadcastingがよくわらないので調べてみた。

2023/12/23に公開

Python

この記事の目的

大学院に出願する際にGNNについての小論文を提出したので、それまでにGNNをちゃんと勉強しておこうということで勉強し始めたのだが、その最初の段階でPyTorchのbroadcastingが出てきてよくわからなかったのでメモ。

検索してqiitaとかzennとかで何かしらの記事出てくるかな？と思っていたが、あまりヒットしなかったので書く。

このページ（How does pytorch broadcasting work?）にbroadcastingについてのドキュメントURLが記載されていたのでそれを読んで理解しようと思う。

紹介されていたドキュメント：

numpy broadcasting rules
PyTorch broadcasting guide.（書いてる途中で大体わかったので、こちらの資料は割愛）

numpy broadcasting rules

このドキュメントの冒頭には

The term broadcasting describes how NumPy treats arrays with different shapes during arithmetic operations. Subject to certain constraints, the smaller array is “broadcast” across the larger array so that they have compatible shapes.

とありますが、"Subject to certain constraints, the smaller array is “broadcast” across the larger array so that they have compatible shapes. "で言ってるように、arrayの形状が異なる場合にはbroadcastingを用いて、ある条件を満たした場合は形状が違っても計算できるようにする機能らしい。

ただ、この機能はメモリの計算効率が悪いため、計算速度を遅延させる可能性がありあまり好ましくない処理のよう。

broadcastingの例

arrayの形状が同じ場合は通常通りに計算できる。

>>> a = np.array([1.0, 2.0, 3.0])
>>> b = np.array([2.0, 2.0, 2.0])
>>> a * b
array([2.,  4.,  6.])

しかし、形状が異なる場合は通常計算できないのだが、broadcastingのおかげで計算できるようになる。この例が最もシンプルなbroadcastingの例である。

この場合は、スカラーと行列の掛け算の計算例である。

>>> a = np.array([1.0, 2.0, 3.0])
>>> b = 2.0
>>> a * b
array([2.,  4.,  6.])

どのような計算をしているのか？

ドキュメント内にわかりやすい画像があったのだが、条件を満たすときは片方のarray（おそらくサイズが小さい方）を変換して計算をできるようにしているらしい。

上で例示したbroadcastingの例では、スカラーのbの方の変数をaのarrayの形状に合わせて計算できるように変換した後に計算して結果を出力している。

In the simplest example of broadcasting, the scalar b is stretched to become an array of same shape as a so the shapes are compatible for element-by-element multiplication.

broadcastingルール

以下のようなルールが適用されるらしく、1,2のどちらかの条件が満たされていれば、broadcastingが適用できるようだ。

条件が満たせない場合は以下のエラーが起きる。
ValueError: operands could not be broadcast together

When operating on two arrays, NumPy compares their shapes element-wise. It >starts with the trailing (i.e. rightmost) dimension and works its way left. >Two dimensions are compatible when

they are equal, or

one of them is 1.

ここは理解がちょっと難しかったのだが、同じ次元の次元数を見た時に、どちらか片方が1次元の場合 or 次元数が一致する場合はbroadcastingが利用できるということだと思う。

例えば、画像のスケール変換で、256x256x3の行列とScaleのスカラーを掛け合わせる時にはbroadcastingが適用できる。

※ 画像スケール変換の例

256x256x3のRGB値の配列があり、画像の各色を異なる値でスケーリングしたい場合、3つの値を持つ1次元配列で画像を乗算することができます。これらの配列の末尾の軸のサイズを放送規則に従って並べると、互換性があることがわかる。

この場合は以下のように、右側から同じ

Image  (3d array): 256 x 256 x 3
Scale  (1d array):             3
Result (3d array): 256 x 256 x 3

絵で表すとこんな感じになり、Scaleの部分がbradcastingにより256 * 256に引き伸ばされて、「256 x 256 x 3」の形状のImage行列との計算を行えるようになるのだと思う。

broadcastingが適用できない例

# 例１
A      (1d array):  3
B      (1d array):  4 # trailing dimensions do not match

# 例２
A      (2d array):      2 x 1
B      (3d array):  8 x 4 x 3 # second from last dimensions mismatched

例１では、片方は3次元で、もう片方は4次元の配列になっているためルールを適用できない。
（Aを4 * 1に拡張すれば計算できるんじゃないか？とか思ったけどルール上できないんでしょう。あ、というかAを4次元にしても、Bを4*3次元の行列に変換して 3*4 x 4*3の行列計算にするという意味のわからないことをするからダメなのか。）

A = [1, 1, 1]
B = [2, 2, 2, 2]

例２では、まぁ見ての通り全然各次元数が合っていないので計算できませんよというわけですね。

ということで大枠broadcastingについてわかったのでここまで。
より詳細についてはこちらのドキュメントを参照されたい。