PyTorch 中的 0/1 背包问题

# 使用 dp_knapsack_sliding_hirschberg——PyTorch 新激活内存预算求解器，峰值 RAM 减少 20 倍 来源：https://jedrzej.maczan.pl/2025_11_21_dp_knapsack_sliding_hirschberg **TL;DR：**如果你使用默认的 `torch._functorch.config.activation_memory_budget_solver=dp`，可以将其改为 `dp_knapsack_sliding_hirschberg`，将内存规划期间的峰值 RAM 使用量降低 20 倍。 #### 是什么？ 当 PyTorch 构建你的 ML 模型的计算图时，它会尝试多种不同的策略来提升速度和内存效率。其中一种技术是：将前向传播和反向传播的联合计算图结合起来，选择哪些操作应该被保存（存储在内存中以便复用），哪些操作应该被重新计算。每个操作执行需要一定时间，存储也需要占用一定内存，而我们可用的内存是有限的。 #### 0/1 背包问题 想象你是一个小偷，刚闯入一个堆满贵重物品的地方。你有一个背包，最多能装 10 公斤。你想最大化偷走物品的价值。理想情况下，你会把所有东西都拿走。但如果你往背包里塞了超过 10 公斤，背包就会破裂，你什么都拿不到。或者，你可以用一些启发式规则来挑选物品（比如只拿小件物品或只拿绿色的），但这样你的方案可能不是最优的——你确实偷到了一些价值，但也许比最优选择能获得的要少。这时，0/1 背包求解器就派上用场了。有不同的算法可以在给定容量（背包最大承重）和给定物品列表（每个物品有价值和成本/重量）的条件下，计算出最优的物品选择。在 PyTorch 中，我们实际做的正是这件事：在 RAM 预算（背包容量）内，选择哪些操作（物品）应该保存到内存中，以最大化运行时的时间节省（价值）。 #### dp_knapsack 目前，PyTorch 中背包求解器的默认实现是 `dp_knapsack`，它使用动态规划方法来求解 0/1 背包问题。一旦你理解了它，算法其实很简单。如果你有过动态规划的经验会很有帮助（我没有）。来看看 Wikipedia 上这个漂亮的 GIF： ByRDSEED (https://commons.wikimedia.org/w/index.php?title=User:RDSEED&action=edit&redlink=1)-Own work，CC BY-SA 4.0 (https://creativecommons.org/licenses/by-sa/4.0)，来源 (https://commons.wikimedia.org/w/index.php?curid=112851814) 左边是 4 个物品的列表。`i` 是索引，`v` 是价值，`w` 是重量。单位无关紧要。容量为 6。这就是为什么右边有 6 列（每列对应总容量的一部分）和 4 行（每个物品一行）。好吧，实际上是 7 列和 5 行，因为用零做了填充，但这只是为了计算方便（每行每列都有父节点，索引算术会简单得多）。你实际上可以不用填充来实现，但我试过之后还是回退到了填充方式，因为索引计算有点烦人。右边的表格称为 DP 表。计算就在这里进行。 #### 构建 DP 表 DP 0/1 背包算法是这样的：取一个物品。在 DP 表中，遍历对应的行，逐列比较你物品的价值与当前 DP 表单元格中的值（行，列；当前列索引就是当前背包的最大重量/容量）。如果你的物品重量大于当前列的索引，跳到下一列。如果你的物品重量等于列索引且物品价值高于当前单元格，用你的物品价值替换该单元格。如果你的物品重量小于列索引且能放入剩余容量，则将其加到单元格的值上（剩余容量 = 当前容量 - 已选物品的容量）。重复直到你遍历完所有物品。此时，你就知道了在重量预算内能获得的最大价值。如果我们只需要测量最大可能价值，到这一步就可以结束了。不幸的是，我们还想知道更多信息。 [](https://jedrzej.maczan.pl/knapsack_no_backtrack.png) #### 回溯 我们需要知道应该选取哪些物品才能得到这个最大价值。所以在 PyTorch 的场景中，就是要知道哪些操作应该保存到内存中，哪些操作应该重新计算。要找出哪些物品对最大价值有贡献（换句话说，哪些物品在构建 DP 表时被选中了），我们需要在 DP 表中反向遍历。我们从最后一行最后一列开始向上。将当前单元格与上方的单元格（同一列，上一行）进行比较。如果它们不同，说明当前物品对最大价值有贡献——因为每一行对应一个物品。如果它们相同，则该物品不会被存入内存，可以重新计算。然后我们向上遍历 DP 表。如果该物品被选中，我们向上移动一行，并将当前单元格的位置向左移动所选物品的重量。如果该物品未被选中，我们向上移动一行，从与之前相同的位置开始。重复直到到达 DP 表的顶部。 [](https://jedrzej.maczan.pl/knapsack_backtrack.png) #### 优化 这个求解器工作正常。它提供了精确解，在大多数情况下你不需要更好的方案。所谓"更好"是指：1. 更快，2. 更省内存。因为再看一下上面的算法和 GIF。它有几个问题。为了得到结果，我们需要分配一个完整的 2D 表，形状为（物品数量，max_weight），这不太理想。**Horace He（Chillee）** (https://github.com/Chillee)——他最初构建了 PyTorch 中内存规划的很大一部分（全部？）——在 DP 背包求解器代码中留下了这条备注 (https://github.com/pytorch/pytorch/blob/1aaedbcfdd5c0615a882eebba3f51b2409162142/torch/_functorch/_activation_checkpointing/knapsack.py#L82)： `# TODO(chilli): I think if needed, this memory can be optimized with sliding # window trick + Hirschberg trick: # https://codeforces.com/blog/entry/47247#comment-316200` 他提出的正是两个优化。 #### 滑动窗口 窗口技巧的意思是，不构建完整的 DP 表，而是在表上滑动，只使用前一行和当前行。移动到下一行时，下一行变为新的当前行，旧的当前行变为前一行。如果只是要在给定容量内从物品列表中计算最大价值，仅这一改进就足够了。但同样地，我们需要知道哪些物品对这个最大价值有贡献。这就是 Horace 建议将滑动窗口与 Hirschberg 技巧结合使用的原因。 [](https://jedrzej.maczan.pl/sliding_window.png) #### Hirschberg 技巧 这是一种分治方法，有点类似于快速排序，因为你会把问题拆分成更小的问题来解决，然后继续拆分，直到解决整个问题。对我们来说，Hirschberg 技巧的主要好处是它消除了回溯。 从实现角度来看，我这里某种程度上做了递归，但没有在栈上压入新帧。我使用一个 Python 列表作为显式的栈（后进先出）。该列表存储物品索引范围和当前可用容量。我们从所有物品和全部容量开始。然后算法是这样的：首先，从栈中弹出一个元素。如果它是一个单独的物品（物品索引范围 end - start = 1），检查物品的重量是否能放入可用容量。如果能，将其加入已保存物品列表；如果不能，加入可重新计算物品列表。在我看来，这部分是 Hirschberg 算法对我们最大的优势——因为你可以看到我们不再需要回溯，因为我们是在算法内部构建已保存和可重新计算的列表。不需要在之后再去构建，不像原始的 0/1 DP 背包算法那样。如果物品范围大于 1，则将物品分成两部分，并为每部分计算一个 DP 表。我们不想构建完整的 DP 表，所以在这里使用滑动窗口技巧，复用当前行和前一行来计算最终的 DP 行。 一旦我们有了左、右两行 DP，将它们相加。找到值最高的位置。那将是我们内存的新分割点。然后我们用这个新计算出的位置将内存一分为二，并使用拆分后的内存和物品再次运行算法。重复直到完成。顺便说一句，可能出现拆分后的某一部分为空的情况。当最优分割点恰好在第一个或最后一个物品时就会发生。这完全正常，我们只需跳过这样的栈帧继续即可。 [](https://jedrzej.maczan.pl/hirschberg_step_by_step.png)这里 (https://codeforces.com/blog/entry/57512) 是 KokiYmgch (https://codeforces.com/profile/KokiYmgch) 在 CodeForces 上对 Hirschberg 算法的详细图文逐步讲解 #### 那么，dp_knapsack_sliding_hirschberg 现在是最佳求解器了吗？ 不是。你看——结合这两个优化，我们可以将 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`。 [](https://jedrzej.maczan.pl/pick_solver.png) #### 如何使用 dp_knapsack_sliding_hirschberg ```python import torch import torch._functorch.config as fconfig fconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschberg ``` 如果你在使用 `dp_knapsack_sliding_hirschberg` 时发现任何 bug，请提交 issue (https://github.com/pytorch/pytorch/issues) 并 @我 @jmaczan (https://github.com/jmaczan)，我会来解决它。 感谢 **Edward Yang** (https://github.com/ezyang) 的代码审查和合并，也感谢你的阅读。