关于 Tensorflow 和 PyTorch 中的自定义操作

Question

我必须实现一个能量函数，称为刚性能量，如本文 here 的公式 7 所示。
能量函数将两个 3D 对象网格作为输入，并返回它们之间的能量。第一个网格是源网格，第二个网格是源网格的变形版本。在粗略的伪代码中，计算过程如下：

遍历源网格中的所有顶点。

1. 对于每个顶点，计算其与其相邻顶点的协方差矩阵。
2. 对计算出的协方差矩阵执行 SVD 并找到顶点的旋转矩阵。
3. 使用计算得到的旋转矩阵、原始网格中的点坐标和变形网格中对应的坐标，计算顶点的能量偏差。

因此，这个能量函数需要我迭代网格中的每个点，并且网格可能有超过 2k 个这样的点。在 TensorFlow 中，有两种方法可以做到这一点。我可以有 2 个形状 (N,3) 的张量，一个代表源点，另一个代表变形网格。

1. 纯粹使用 Tensorflow 张量。也就是说，使用tf.gather 迭代上述张量的元素，并仅使用现有的 TF 操作对每个点执行计算。这种方法，会非常慢。我之前曾尝试定义迭代超过 1000 个点的损失函数，但图形构造本身需要太多时间才能实用。
2. 按照 TF 文档 here 中的说明添加新的 TF OP。这涉及在 CPP（和 Cuda，用于 GPU 支持）中编写函数，并将新的 OP 注册到 TF​​。

第一种方法很容易编写，但速度慢得不切实际。第二种方法写起来很痛苦。

我使用 TF 已经 3 年了，之前从未使用过 PyTorch，但此时我正在考虑切换到它，如果它为这种情况提供了更好的替代方案。

PyTorch 是否有一种方法可以轻松实现此类损失函数，并且执行速度与在 GPU 上一样快。即，我自己编写在 GPU 上运行的损失函数的 Pythonic 方式，我没有任何 C 或 Cuda 代码？

Answer 1

据我了解，您本质上是在询问此操作是否可以矢量化。答案是否定的，至少不完全，因为 PyTorch 中的 svd 实现不是矢量化的。

如果您展示了 tensorflow 实现，它将有助于理解您的起点。我不知道找到顶点的旋转矩阵是什么意思，但我想这可以向量化。这意味着 svd 是唯一的非向量化操作，您也许可以只编写一个自定义 OP，即向量化 svd - 这可能很容易，因为它相当于在循环中调用一些库例程在 C++ 中。

我看到的两个可能的问题来源是

1. 如果等式 7 中 N(i) 的邻域可以具有显着不同的大小（这意味着协方差矩阵具有不同的大小并且向量化将需要一些肮脏的技巧）
2. 处理网格和邻域的一般问题可能很困难。这是不规则网格的固有属性，但 PyTorch 支持稀疏矩阵和专用包 torch_geometry，这至少有帮助。