PyTorch v1.12 以降では、macOS において Apple Silicon あるいは AMD の GPU を使ったアクセラレーションが可能になっているらしい。
バックエンドの名称は Metal Performance Shaders (MPS) という。
意外と簡単に使えるようなので、今回は手元の Mac で試してみた。
使った環境は次のとおり。
GPU が 19 コアの Apple M2 Pro を積んだ Mac mini を使用している。
$ sw_vers
ProductName: macOS
ProductVersion: 14.4.1
BuildVersion: 23E224
$ sysctl machdep.cpu.brand_string
machdep.cpu.brand_string: Apple M2 Pro
$ pip list | grep -i torch
torch 2.2.1
$ python -V
Python 3.10.14
もくじ
下準備
あらかじめ、必要なパッケージをインストールする。
特に意識しなくても MPS バックエンドが有効なバイナリが入る。
$ pip install torch tqdm numpy
インストールできたら Python のインタプリタを起動する。
$ python
そして、PyTorch のパッケージをインポートしておく。
>>> import torch
MPS バックエンドを使ってみる
MPS バックエンドが有効かどうかは以下のコードで確認できる。
True が返ってくれば利用できる状態にある。
>>> torch.backends.mps.is_available()
True
使い方は CUDA バックエンドと変わらない。
テンソルやモデルを .to() メソッドで転送するだけ。
このとき、引数に "mps" を指定すれば良い。
>>> x = torch.randn(2, 3, 4).to("mps")
>>> x.shape
torch.Size([2, 3, 4])
>>> x.device
device(type='mps', index=0)
ちゃんと転送できた。
簡単にベンチマークしてみる
続いては、どれくらいパフォーマンスが出るのか気になるので簡単にベンチマークしてみる。
PyTorch のベンチマークのページ 1 を参考に、以下のようなコードを用意した。
いくつかのサイズやスレッド数の組み合わせで、行列の積や和を計算している。
from itertools import product
from tqdm import tqdm
import torch
import torch.utils.benchmark as benchmark
def device():
"""環境毎に利用できるアクセラレータを返す"""
if torch.backends.mps.is_available():
return "mps"
if torch.cuda.is_available():
return "cuda"
return "cpu"
def batched_dot_mul_sum(a, b):
"""mul -> sum"""
return a.mul(b).sum(-1)
def batched_dot_bmm(a, b):
"""bmm -> flatten"""
a = a.reshape(-1, 1, a.shape[-1])
b = b.reshape(-1, b.shape[-1], 1)
return torch.bmm(a, b).flatten(-3)
DEVICE = device()
print(f"device: {DEVICE}")
results = []
sizes = [1, 64, 1024, 10000]
for b, n in tqdm(list(product(sizes, sizes))):
label = "Batched dot"
sub_label = f"[{b}, {n}]"
x = torch.ones((b, n)).to(DEVICE)
for num_threads in [1, 4, 16, 32]:
results.append(benchmark.Timer(
stmt="batched_dot_mul_sum(x, x)",
setup="from __main__ import batched_dot_mul_sum",
globals={"x": x},
num_threads=num_threads,
label=label,
sub_label=sub_label,
description="mul/sum",
).blocked_autorange(min_run_time=1))
results.append(benchmark.Timer(
stmt="batched_dot_bmm(x, x)",
setup="from __main__ import batched_dot_bmm",
globals={"x": x},
num_threads=num_threads,
label=label,
sub_label=sub_label,
description="bmm",
).blocked_autorange(min_run_time=1))
compare = benchmark.Compare(results)
compare.print()
Apple M2 Pro (GPU 19C)
実際に、上記を実行してみよう。
まずは Apple M2 Pro の環境から。
$ python bench.py
device: mps
100%|███████████████████████████████████████████| 16/16 [02:12<00:00, 8.27s/it]
[-------------- Batched dot --------------]
| mul/sum | bmm
1 threads: --------------------------------
[1, 1] | 49.9 | 30.8
[1, 64] | 48.7 | 30.1
[1, 1024] | 48.8 | 30.0
[1, 10000] | 51.5 | 30.1
[64, 1] | 50.1 | 30.3
[64, 64] | 49.1 | 30.2
[64, 1024] | 54.9 | 30.1
[64, 10000] | 58.0 | 30.0
[1024, 1] | 49.9 | 30.0
[1024, 64] | 55.5 | 30.4
[1024, 1024] | 54.9 | 30.0
[1024, 10000] | 400.2 | 90.0
[10000, 1] | 53.7 | 30.5
[10000, 64] | 56.0 | 31.0
[10000, 1024] | 271.2 | 107.0
[10000, 10000] | 6594.7 | 31.2
4 threads: --------------------------------
[1, 1] | 52.1 | 31.5
[1, 64] | 50.6 | 31.3
[1, 1024] | 50.5 | 30.5
[1, 10000] | 53.2 | 31.3
[64, 1] | 52.7 | 31.3
[64, 64] | 51.2 | 30.3
[64, 1024] | 56.7 | 30.5
[64, 10000] | 59.6 | 30.7
[1024, 1] | 51.5 | 30.6
[1024, 64] | 56.6 | 30.7
[1024, 1024] | 57.1 | 30.7
[1024, 10000] | 64.5 | 204.3
[10000, 1] | 55.3 | 35.1
[10000, 64] | 58.0 | 34.4
[10000, 1024] | 590.8 | 223.3
[10000, 10000] | 32409.0 | 1498.3
16 threads: -------------------------------
[1, 1] | 51.6 | 30.8
[1, 64] | 51.1 | 30.4
[1, 1024] | 50.6 | 30.4
[1, 10000] | 53.7 | 30.7
[64, 1] | 51.7 | 30.6
[64, 64] | 50.4 | 30.4
[64, 1024] | 57.1 | 30.7
[64, 10000] | 59.5 | 30.5
[1024, 1] | 51.2 | 30.3
[1024, 64] | 56.3 | 30.8
[1024, 1024] | 57.3 | 31.0
[1024, 10000] | 60.3 | 106.8
[10000, 1] | 54.9 | 34.9
[10000, 64] | 57.2 | 34.5
[10000, 1024] | 400.3 | 220.7
[10000, 10000] | 32418.2 | 1503.2
32 threads: -------------------------------
[1, 1] | 51.1 | 30.6
[1, 64] | 50.4 | 30.6
[1, 1024] | 50.7 | 30.5
[1, 10000] | 53.0 | 30.5
[64, 1] | 51.8 | 30.7
[64, 64] | 50.4 | 30.2
[64, 1024] | 56.7 | 30.6
[64, 10000] | 59.3 | 30.5
[1024, 1] | 51.3 | 30.6
[1024, 64] | 56.6 | 34.5
[1024, 1024] | 57.8 | 33.5
[1024, 10000] | 447.3 | 202.6
[10000, 1] | 54.3 | 35.3
[10000, 64] | 57.0 | 34.5
[10000, 1024] | 591.2 | 219.7
[10000, 10000] | 32443.3 | 1493.3
Times are in microseconds (us).
NVIDIA GeForce RTX 3060
さきほどの結果は、もちろん CPU よりは全然速い。
とはいえ、他の GPU などに比べてどれくらい速いのかイメージしにくい。
そこで、厳密な比較にはならないものの RTX 3060 を積んだ Linux の環境でも実行してみる。
$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description: Ubuntu 20.04.6 LTS
Release: 20.04
Codename: focal
$ pip list | grep -i torch
torch 2.2.2
$ python -V
Python 3.10.14
先ほどのコードを実行する。
$ python bench.py
device: cuda
100%|███████████████████████████████████████████| 16/16 [03:08<00:00, 11.81s/it]
[-------------- Batched dot --------------]
| mul/sum | bmm
1 threads: --------------------------------
[1, 1] | 6.5 | 6.8
[1, 64] | 6.5 | 6.8
[1, 1024] | 6.4 | 7.8
[1, 10000] | 6.3 | 7.8
[64, 1] | 6.3 | 6.6
[64, 64] | 6.4 | 6.8
[64, 1024] | 6.6 | 6.8
[64, 10000] | 25.1 | 10.2
[1024, 1] | 6.4 | 6.7
[1024, 64] | 6.3 | 6.7
[1024, 1024] | 40.2 | 15.6
[1024, 10000] | 375.3 | 179.1
[10000, 1] | 6.3 | 32.6
[10000, 64] | 29.2 | 34.9
[10000, 1024] | 374.7 | 123.5
[10000, 10000] | 3603.7 | 1672.6
4 threads: --------------------------------
[1, 1] | 6.5 | 6.9
[1, 64] | 6.5 | 6.9
[1, 1024] | 6.4 | 7.8
[1, 10000] | 6.4 | 7.8
[64, 1] | 6.4 | 6.6
[64, 64] | 6.5 | 6.8
[64, 1024] | 6.6 | 6.9
[64, 10000] | 25.1 | 10.2
[1024, 1] | 6.3 | 6.7
[1024, 64] | 6.4 | 6.7
[1024, 1024] | 40.2 | 15.6
[1024, 10000] | 375.3 | 179.1
[10000, 1] | 6.3 | 32.6
[10000, 64] | 29.2 | 34.9
[10000, 1024] | 374.9 | 123.5
[10000, 10000] | 3602.4 | 1672.5
16 threads: -------------------------------
[1, 1] | 6.5 | 6.9
[1, 64] | 6.5 | 6.7
[1, 1024] | 6.5 | 7.9
[1, 10000] | 6.3 | 7.8
[64, 1] | 6.3 | 6.6
[64, 64] | 6.4 | 6.8
[64, 1024] | 6.5 | 6.9
[64, 10000] | 25.1 | 10.2
[1024, 1] | 6.3 | 6.7
[1024, 64] | 6.4 | 6.7
[1024, 1024] | 40.3 | 15.6
[1024, 10000] | 375.3 | 179.1
[10000, 1] | 6.4 | 32.6
[10000, 64] | 29.2 | 34.9
[10000, 1024] | 374.9 | 123.5
[10000, 10000] | 3604.9 | 1672.4
32 threads: -------------------------------
[1, 1] | 6.6 | 6.9
[1, 64] | 6.4 | 6.8
[1, 1024] | 6.5 | 7.9
[1, 10000] | 6.4 | 7.8
[64, 1] | 6.3 | 6.7
[64, 64] | 6.6 | 6.8
[64, 1024] | 6.6 | 6.9
[64, 10000] | 25.1 | 10.3
[1024, 1] | 6.4 | 6.8
[1024, 64] | 6.3 | 6.8
[1024, 1024] | 40.2 | 15.6
[1024, 10000] | 375.1 | 179.2
[10000, 1] | 6.4 | 32.6
[10000, 64] | 29.2 | 34.9
[10000, 1024] | 374.9 | 123.5
[10000, 10000] | 3604.6 | 1672.4
Times are in microseconds (us).
こちらの環境の方が多くの場合に 2 ~ 10 倍程度速いことがわかる。
ただし、一部サイズの大きな bmm を使った演算に関しては、むしろ Apple Silicon の方が速いようだ。
また、消費電力は RTX 3060 の方が 20 倍近く大きい 2。
まとめ
Apple Silicon の GPU は、そこまで速くないにしてもワットパフォーマンスには優れている。
また、CPU に比べればずっと速いので PyTorch で気軽に使えるのはありがたい。
参考
developer.apple.com
pytorch.org
pytorch.org
