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

Question

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

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

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

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

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

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

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

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

Answer 1

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

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

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

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