关于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支持）编写函数，以及用TF注册新的OP。

第一种方法很容易编写，但实际上却很慢。第二种方法很难写。

我已经使用TF三年了，以前从未使用过PyTorch，但是目前，我正在考虑切换到它，如果它可以为此类情况提供更好的选择。

PyTorch是否可以轻松实现 这样的损失函数，并且执行速度与在GPU上一样快。即一种编写在GPU上运行我自己的损失函数的Python方法，而我却不需要任何C或Cuda代码？

Answer 1

据我了解，您实际上是在问是否可以对该操作进行向量化。答案是否定的，至少是不完全的，因为PyTorch中的svd实现不是矢量化的。

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

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

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