PyTorchにおける0/1ナップサック問題

0/1ナップサック問題 in PyTorch

TL;DR: torch._functorch.config.activation_memory_budget_solver のデフォルト設定 dp を dp_knapsack_sliding_hirschberg に変更すると、メモリ計画時のピークRAM使用量を20分の1に削減できます。

［画像: ナップサック問題の概念図］

何か？

PyTorchが機械学習モデルの計算グラフを構築する際、高速化とメモリ効率化のために様々な手法を試みます。その一つが、フォワードパスとバックワードパスの計算グラフを統合し、どの演算を保存（メモリに格納して再利用）し、どの演算を再計算するかを選択する技術です。各演算は実行に時間がかかり、保存するとメモリを消費しますが、利用できるメモリ量には限りがあります。

0/1ナップサック問題

あなたが泥棒で、貴重品がたくさんある場所に侵入したと想像してください。あなたは最大10kgまで運べるナップサックを持っています。盗むものの価値を最大化したいと考えています。理想的には、すべて盗みたいでしょう。しかし、ナップサックに10kg以上のものを詰め込むと、壊れてしまい、何も持ち帰れなくなります。あるいは、小さなものだけを盗む、緑色のものだけを盗むといったヒューリスティックな方法でアイテムを選ぶこともできますが、その場合、解は最適ではないかもしれません。つまり、ある程度の価値は得られますが、最適な選択をした場合に得られたであろう価値よりも少ないかもしれません。そこで、0/1ナップサックソルバーが登場します。これらは、与えられた容量（ナップサックに詰め込める最大重量）と、各アイテムに価値とコスト（重量）があるアイテムのリストに対して、最適なアイテムの選択を計算する様々なアルゴリズムです。PyTorchでは、これを以下のように利用しています。RAMの予算（ナップサックの容量）内で、実行時間の節約（価値）を最大化するために、どの演算（アイテム）をメモリに保存すべきかを選択します。

［画像: 泥棒がナップサックにアイテムを詰めている様子］

dp_knapsack

現在、PyTorchにおけるナップサックソルバーのデフォルト実装は dp_knapsack で、これは0/1ナップサック問題を解くための動的計画法（DP）アプローチです。このアルゴリズムは、理解すれば実際にはシンプルです。動的計画法に事前に触れた経験があると、大いに役立ちます（私はありませんでした）。Wikipediaから引用したこの素晴らしいGIFを見てみましょう。

［画像: 動的計画法によるナップサック問題のDPテーブル構築のGIF］ By RDSEED - Own work, CC BY-SA 4.0, Link

左側には4つのアイテムのリストがあります。iはインデックス、vは価値、wは重量です。単位は関係ありません。容量は6です。そのため、右側には6つの列（各列は総容量の一部）と4つの行（アイテムごと）があります。実際には、ゼロパディングのために7列と5行ありますが、これは計算を容易にするためのものです（各行と列には親があり、インデックス演算がはるかに簡単になります）。パディングなしでも実行できますが、試した結果、インデックス演算が少し面倒だったので、結局パディングを使用しました。右側の表はDPテーブルと呼ばれ、計算が行われる場所です。

DPテーブルの構築

DP 0/1ナップサックアルゴリズムは次のようになります。アイテムを取得します。DPテーブル内で、対応する行を列ごとに見ていき、アイテムの価値をDPテーブルの現在のセル（行、列；現在の列インデックスはナップサックの現在の最大重量（容量））と比較します。アイテムの重量が現在の列のインデックスよりも大きい場合は、次の列に進みます。アイテムの重量が列のインデックスと等しく、アイテムの価値が現在のセルよりも高い場合は、セルをアイテムの価値で置き換えます。アイテムの重量が列のインデックスよりも小さく、余剰容量に収まる場合は、セルの値（空き容量は現在の容量から選択済みアイテムの容量を引いたもの）にそれを加算します。すべてのアイテムを繰り返し処理するまで繰り返します。この時点で、重量予算内で得られる最大価値を知ることができます。もし目標が最大可能値を測定することであれば、ここで終了します。残念ながら、私たちはもっと知りたいのです。

［画像: バックトラッキングなしのナップサック問題のDPテーブル］

バックトラッキング

この最大価値を得るためにどのアイテムを取るべきかを知る必要があります。PyTorchの場合、どの演算をメモリに保存し、どの演算を再計算すべきかを知る必要があります。最大価値に貢献するアイテム（つまり、DPテーブル構築中に選択されたアイテム）を見つけるには、DPテーブルを逆方向にたどる必要があります。最後の行と最後の列から上に移動します。セルをその上のセル（同じ列、前の行）と比較します。それらが異なる場合、現在のアイテムは最大価値に貢献しています。なぜなら、各行は単一のアイテムに対応しているからです。それらが等しい場合、このアイテムはメモリに保存されず、再計算可能になります。次に、DPテーブルを上にたどります。アイテムが選択された場合、1行上に移動し、現在のセルの位置を、選択されたアイテムの重量だけ左に移動します。アイテムが選択されなかった場合、1行上に移動し、以前と同じ位置から開始します。DPテーブルの最上部に到達するまで繰り返します。

［画像: バックトラッキングありのナップサック問題のDPテーブル］

最適化

このソルバーは正しく動作します。正確な解を提供し、ほとんどの場合、それ以上のものは必要ありません。より良いものとは、1. より高速で、2. よりメモリを消費しないことを意味します。アルゴリズムと上記のGIFをもう一度見てください。いくつかの問題点があります。結果を得るために、形状が（アイテム数、最大重量）の完全な2Dテーブルを割り当てますが、これはあまり良くありません。PyTorchのメモリ計画の大部分（すべて？）を元々構築したHorace He (Chillee)は、DPナップサックソルバーのコードに次のようなメモを残しています。

# TODO(chilli): 必要であれば、このメモリはスライディングウィンドウトリックとハーシュバーグトリックで最適化できると思います：# https://codeforces.com/blog/entry/47247?#comment-316200

彼が提案しているのは2つの最適化です。

スライディングウィンドウ

ウィンドウトリックとは、DPテーブルを構築する代わりに、テーブル上をスライドし、前の行と現在の行のみを使用することを意味します。次の行に移動する際、新しい行が現在の行になり、古い現在の行が前の行になります。この改善だけでも、リストからアイテムの最大価値を計算したいだけであれば十分です。しかし、再び、どのアイテムがこの最大価値に貢献するかを知る必要があります。そのため、Horaceはスライディングウィンドウとハーシュバーグトリックを組み合わせて使用することを提案しました。

［画像: スライディングウィンドウの概念図］

ハーシュバーグトリック

これは分割統治アプローチで、クイックソートに似ています。問題を小さな問題に分割して解決し、分割を続けて全体の問題を解決します。ハーシュバーグトリックの主な利点は、バックトラッキングを排除できることです。

実装の観点からは、スタックに新しいフレームをプッシュすることなく、再帰のようなことを行っています。Pythonリストを明示的なスタック（後入れ先出し）として使用しています。リストには、アイテムインデックスの範囲と、現在利用可能な容量が格納されます。すべてのアクティビティとすべての容量から開始します。次に、アルゴリズムは次のように進行します。まず、スタックからアイテムをポップします。それが単一のアイテム（アイテムインデックスの範囲 end - start = 1）である場合、アイテムの重量が利用可能な容量に収まるかどうかを確認します。収まる場合は、保存済みアイテムに含め、収まらない場合は、再計算可能アイテムに含めます。この部分は、ハーシュバーグアルゴリズムにとって最大の強みだと思います。なぜなら、アルゴリズム内で保存済みリストと再計算可能リストを構築するため、後でバックトラッキングする必要がないことがわかるからです。元の0/1 DPナップサックアルゴリズムのように、後で構築する必要はありません。アイテムの範囲が1より大きい場合、アイテムを2つの部分に分割し、それぞれについてDPテーブルを計算します。完全なDPテーブルを構築したくないため、ここでスライディングウィンドウトリックを使用し、現在のDP行と前のDP行を再利用して最終的なDP行を計算します。

両方の左側と右側のDP行が得られたら、それらを合計します。最も高い値の位置を見つけます。これがメモリの新しい分割点になります。次に、この新しく計算された位置を使用してメモリを2つに分割し、メモリとアイテムの分割の両方を使用してアルゴリズムを再度実行します。完了するまで繰り返します。ちなみに、分割された部分のいずれかが空になる場合があります。これは、最適な分割が最初のアイテムまたは最後のアイテムのいずれかである場合に発生する可能性があります。これは完全に正当であり、そのようなスタックフレームをスキップして続行するだけです。

［画像: ハーシュバーグアルゴリズムのステップバイステップの説明］こちらは、CodeForcesのKokiYmgchによるハーシュバーグアルゴリズムの描画付きの素晴らしいステップバイステップの説明です。

では、dp_knapsack_sliding_hirschberg は現在のTHEソルバーなのでしょうか？

いいえ。これらの2つの最適化を組み合わせることで、DPテーブルのサイズを形状（アイテム数、最大重量）から形状（2、最大重量）に減らすことができます。これにより、ピークメモリを効果的に20倍削減できます。私の64GB RAMを搭載したマシンでは、dp_knapsack は100アイテムの問題でクラッシュしますが、dp_knapsack_sliding_hirschberg はこのRAM予算内で2000アイテムの問題を処理できます。実行速度も約37%向上しましたが、これは厳密には測定していません。とはいえ、SciPyの使用を気にしないのであれば、おそらく ilp_knapsack を使用すべきです。これは、PyTorchの既存のDPナップサック実装よりもはるかに高速です。貪欲アルゴリズムよりも少し遅いですが、それでも非常に高速であり、貪欲アルゴリズムとは異なり、両方のDPソルバーと同様に、正確な解を提供します。

あなたに合ったソルバーの選び方

速度を正しさよりも重視する場合は、greedy_knapsack を使用してください。

それ以外で、プロジェクトにSciPyを導入しても問題ない場合は、ilp_knapsack を使用してください。

それ以外で、パッケージマネージャーからPyTorchをインストールしている場合は、何もしないでください。PyTorchは dp_knapsack を使用します。これはデフォルトのソルバーであり、何もする必要はありません。

それ以外の場合は、dp_knapsack_sliding_hirschberg を使用してください。

［画像: ソルバー選択ガイド］

dp_knapsack_sliding_hirschberg の使い方

import torch
import torch._functorch.config as fconfig

fconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschberg

dp_knapsack_sliding_hirschberg を使用中にバグを見つけた場合は、問題を登録してください。私（@jmaczan）にメンションしていただければ、修正に取り組みます。

コードレビューとマージをしてくれたEdward Yangに感謝します。そして、読んでくれたあなたにも感謝します。