我必须实现一个能量函数,称为刚性能量,如本文的等式 7 所示 here .
能量函数将两个 3D 对象网格作为输入,并返回它们之间的能量。第一个网格是源网格,第二个网格是源网格的变形版本。在粗略的伪代码中,计算将如下所示:
遍历源网格中的所有顶点。
- 对于每个顶点,计算其与其相邻顶点的协方差矩阵。
- 对计算出的协方差矩阵执行 SVD 并找到顶点的旋转矩阵。
- 使用计算出的旋转矩阵、原始网格中的点坐标和变形后网格中的相应坐标,计算顶点的能量偏差。
因此这个能量函数需要我遍历网格中的每个点,而网格可能有超过 2k 个这样的点。在 Tensorflow 中,有两种方法可以做到这一点。我可以有 2 个形状为 (N,3) 的张量,一个代表源点,另一个代表变形网格。
- 完全使用 Tensorflow 张量来完成。也就是说,使用
tf.gather
迭代上述张量的元素,并仅使用现有的 TF 操作对每个点执行计算。这种方法,会非常慢。我之前尝试定义迭代超过 1000 个点的损失函数,并且图形构建本身需要太多时间而不实用。 - 按照 TF 文档中的说明添加新的 TF OP here .这涉及在 CPP(和 Cuda,用于 GPU 支持)中编写函数,并使用 TF 注册新的 OP。
第一种方法很容易写,但速度慢得不切实际。第二种方法写起来很痛苦。
我已经使用 TF 3 年了,以前从未使用过 PyTorch,但现在我正在考虑切换到它,如果它能为此类情况提供更好的替代方案。
PyTorch 是否有一种方法可以轻松实现此类损失函数,并且的执行速度与在 GPU 上一样快。即,一种编写我自己的在 GPU 上运行的损失函数的 pythonic 方式,我没有任何 C 或 Cuda 代码?
最佳答案
据我了解,您实质上是在询问此操作是否可以矢量化。答案是否定的,至少不完全是,因为svd PyTorch 中的实现未矢量化。
如果您展示了 tensorflow 实现,将有助于理解您的起点。我不知道你所说的找到顶点的旋转矩阵是什么意思,但我猜这可以被矢量化。这意味着 svd 是唯一的非矢量化操作,您也许可以只编写一个自定义 OP,即矢量化 svd - 这可能很容易,因为它相当于在循环中调用一些库例程在 C++ 中。
我看到的两个可能的问题来源是
- 如果等式 7 中
N(i)
的邻域大小可能明显不同(这意味着协方差矩阵大小不同,向量化需要一些卑鄙的技巧) - 处理网格和邻域的一般问题可能很困难。这是不规则网格的固有属性,但 PyTorch 支持稀疏矩阵和专用包 torch_geometry ,这至少有帮助。
关于python - 关于 Tensorflow 和 PyTorch 中的自定义操作,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/54473620/