我要实现的能量函数，称为刚性能源，如本文的公式7这里。
能量函数将两个 3D 对象网格作为输入，并返回它们之间的能量。第一个网格是源网格，第二个网格是源网格的变形版本。在粗略的伪代码中，计算过程如下：
迭代源网格中的所有顶点。
因此，这个能量函数要求我迭代网格中的每个点，并且网格可能有超过 2k 个这样的点。在 Tensorflow 中，有两种方法可以做到这一点。我可以有 2 个形状 (N,3) 的张量，一个代表源点，另一个代表变形网格。
tf.gather使用现有的 TF 操作迭代上述张量的元素并在每个点上执行计算。这种方法，会非常慢。我之前曾尝试定义迭代超过 1000 个点的损失函数，而图构建本身需要太多时间而无法实用。第一种方法很容易编写，但速度很慢。第二种方法写起来很痛苦。
我已经使用 TF 3 年了，之前从未使用过 PyTorch，但此时我正在考虑切换到它，如果它为这种情况提供了更好的替代方案。
PyTorch 是否有一种方法可以轻松实现此类损失函数，并且执行速度与在 GPU 上一样快。即，编写我自己的在 GPU 上运行的损失函数的 pythonic 方式，我没有任何 C 或 Cuda 代码？
据我了解，您本质上是在问这个操作是否可以矢量化。答案是否定的，至少不完全是，因为PyTorch 中的svd实现不是矢量化的。
如果您展示了 tensorflow 实现，它将有助于理解您的起点。我不知道你通过找到顶点的旋转矩阵是什么意思，但我猜这可以被矢量化。这意味着 svd 是唯一的非矢量化操作，您也许可以只编写一个自定义 OP，即矢量化 svd - 这可能很容易，因为它相当于在循环中调用一些库例程在 C++ 中。
我看到的两个可能的问题来源是
N(i)等式 7 中的邻域的大小可能有显着差异（这意味着协方差矩阵的大小不同，并且矢量化需要一些肮脏的技巧）本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句