PyTorchとNumpyでsortのベンチマーク

PyTorchとNumpyでSortを簡単に計ってみた。

PyTorchでは、sortアルゴリズムとして、クイックソート(quick sort)と基数ソート(radix sort)を使っている。クイックソートは、PyTorchベタに実装されているものであり、基数ソートはCUDAのライブラリthrust::sortで実装されている。また、numpyは基本はクイックソートであるが、それ以外のアルゴリズムも実装されている。

以下のコードで測定したので、Pythonから見て純粋な計算時間である。ここでは、2次元テンソル(100x100, 1000x1000, 10000x10000)て調べた。

import torch
import numpy as np
a = torch.randn(100,100)
b = torch.randn(1000,1000)
c = torch.randn(10000,10000)
%timeit torch.sort(a,1)
%timeit torch.sort(b,1)
%timeit torch.sort(c,1)
a = a.cuda()
b = b.cuda()
c = c.cuda()
%timeit torch.sort(a,1)
%timeit torch.sort(b,1)
%timeit torch.sort(c,1)
d = np.random.randn(100,100)
%timeit np.sort(d,1)
e = np.random.randn(1000,1000)
%timeit np.sort(e,1)
f = np.random.randn(10000,10000)
%timeit np.sort(f,1)

上のデータをpandasでプロットする。CUDA(Tesla T4)のほうが、2桁近く性能改善している。なお、K80でも2倍程度時間がかかるだけでそれでも早いことがわかる。また、CPUでの単純な比較では、numpyがpytorchより2割程度性能が良い。なお、xは上記の要素数。yはマイクロ秒(usec)である。

上の図を描画するスクリプトは以下である。なお、日本語化したい場合、Japanize-matplotlibを使えば、2行追加で日本語化(IPAゴシックフォント化)できる。

%matplotlib inline
import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({'x': [100, 1000, 10000], 'torch_cpu': [403, 57200, 7190000],
                    'torch_cuda_t4': [25, 782, 401000], 'torch_cuda_k80' : [37, 1640, 694000],
                    'numpy': [289, 44000, 5850000]})
fig, ax = plt.subplots(figsize=(5, 3))
df.set_index('x').plot(
    ax=ax,
    logy=True
)
ax.set_title('Sort elapstime (usec)')
# ax.set_xlabel('')
ax.set_ylabel('y')
plt.show()

参考資料

• ソースコード
  • at::sort
    • THCTensor_(sort) (CUDAの場合)
      • PyTorchのCUDAでは、スライス時一定サイズ(2048要素)以上のソートは、thrust::sortを用いている。
        • thrust sorting
    • THTensor_(sort)(CPUの場合)
• uehara1414/japanize-matplotlib matplotlibの日本語化
  • matplotlib日本語対応　日本語化matplotlibの使い方
  • (余談)seabornを使う場合、sns.set()を実行すると、フォント設定が上書きされてしまう。このため、そのあとにimport japanize-matplotlibとする必要がある。