高速微分ソートアルゴリズム、カスタムC++、CUDAのパッケージで、パフォーマンスが向上しました。

[[390293]]

昨年の初め、Google は論文「Fast Differentiable Sorting and Ranking」で、O(nlogn) の時間計算量と O(n) の空間計算量を持つ初の微分可能ソート アルゴリズムを発表しました。これは、既存の方法よりも桁違いに高速です。

最近、誰かが GitHub でプロジェクトをオープンソース化しました。これは、ソフトウェア パッケージの形式で高速な微分ソートとランキングを実装したものです。このプロジェクトは数日間オンラインになり、300 以上のスターを獲得しました。

• プロジェクトアドレス: https://github.com/teddykoker/torchsort

• 高速微分ソートおよびランキングの論文アドレス: https://arxiv.org/pdf/2002.08871.pdf

トーチソート

Torchsort は、Blondel らによって提案された高速微分ソートおよびランキングを実装し、純粋な PyTorch で実装されています。コードの大部分は、プロジェクト「google-research/fast-soft-sort」のオリジナルの Numpy 実装からコピーされており、高速パフォーマンスのためにカスタム C++ および CUDA カーネルが使用されています。

Torchsort のインストールは非常に簡単です。一般的に使用されている pip インストール方法を使用できます。インストール コードは次のとおりです。

 pip インストール torchsort

CUDA 拡張機能をビルドする場合は、CUDA ツールチェーンをインストールする必要があります。 CUDA 以外のランタイム環境で docker などのアプリケーションをビルドする場合は、インストール前に環境変数「TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"」をエクスポートする必要があります。

使い方

torchsort には、soft_rank と soft_sort という 2 つの関数があり、それぞれに、regularization (l2 または kl) (正規化関数) とregularization_strength (スカラー値) というパラメーターがあります。それぞれは、正規化の強度に応じて精度を変えながら、2D テンソルの最後の次元をソートします。

輸入トーチ
トーチソートをインポートする
x = torch.tensor([[ 8 , 0 , 5 , 3 , 2 , 1 , 6 , 7 , 9 ]])
 torchsort.soft_sort(x, 正規化の強度= 1.0 )
 # テンソル([[ 0.5556 , 1.5556 , 2.5556 , 3.5556 , 4.5556 , 5.5556 , 6.5556 , 7.5556 , 8.5556 ]])
 torchsort.soft_sort(x, 正規化の強度= 0.1 )
 # テンソル([[- 0 ., 1 ., 2 ., 3 ., 5 ., 6 ., 7 ., 8 ., 9 .]])
トーチソート.ソフトランク(x)
 # テンソル([[ 8 ., 1 ., 5 ., 4 ., 3 ., 2 ., 6 ., 7 ., 9 .]])

どちらの操作も完全に微分可能であり、次のように CPU または GPU に実装できます。

 x = torch.tensor([[ 8 ., 0 ., 5 ., 3 ., 2 ., 1 ., 6 ., 7 ., 9 .]], requires_grad=True).cuda()
 y = トーチソート.soft_sort(x)
 torch.autograd.grad(y[ 0,0 ],x) の関数
# (テンソル([[ 0.1111 , 0.1111 , 0.1111 , 0.1111 , 0.1111 , 0.1111 , 0.1111 , 0.1111 ]],
 # デバイス = 'cuda:0' ),)

表示例

スピアマン順位係数は、2 つの変数間の単調な相関を測定するのに非常に役立つ指標です。 Torchsort を使用して微分可能なスピアマン順位係数関数を作成し、そのメトリックに対してモデルを直接最適化することができます。

輸入トーチ
トーチソートをインポートする
定義spearmanr(予測、ターゲット、**kw):
    pred = torchsort.soft_rank(pred, **kw)
    ターゲット = torchsort.soft_rank(ターゲット、**kw)
    pred = pred - pred.mean()
    pred = pred / pred.norm()
    ターゲット = ターゲット - ターゲット.mean()
    ターゲット = ターゲット / ターゲット.norm()
戻り値(予測 * ターゲット).sum()
 pred = torch.tensor([[ 1 ., 2 ., 3 ., 4 ., 5 .]], requires_grad=True)
ターゲット = torch.tensor([[ 5 ., 6 ., 7 ., 8 ., 7 .]])
スピアマン = スピアマンr(予測、ターゲット)
 # テンソル( 0.8321 )
 torch.autograd.grad(スピアマン、予測)
 # (テンソル([[- 5 .5470e- 02 , 2 .9802e- 09 , 5 .5470e- 02 , 1 .1094e- 01 , - 1 .1094e- 01 ]]),)

ベンチマーク

torchsort と fast_soft_sort はどちらも時間計算量が O(n log n) で、組み込みの torch.sort と比較してオーバーヘッドが若干追加されます。 Numba JIT のバッチ サイズが 1 の場合 (左の図を参照)、fast_soft_sort のフォワード パスは Torchsort CPU カーネルとほぼ同じパフォーマンスを発揮しますが、バックワード パスは依然として一部の Python コードに依存しているため、パフォーマンスが大幅に低下します。

さらに、torchsort カーネルはバッチ処理をサポートしており、バッチ サイズが大きくなるにつれて fast_soft_sort よりも優れたパフォーマンスが得られます。

トーチソート CUDA カーネルは、2000 未満のシーケンス長では優れたパフォーマンスを発揮し、非常に大きなバッチに拡張できます。将来的には、CUDA カーネルがさらに最適化され、組み込みの torch.sort に近いパフォーマンスが実現される可能性があります。