我要實現的能量函數，稱為剛性能源，如本文的公式7這里。能量函數將兩個 3D 對象網格作為輸入，并返回它們之間的能量。第一個網格是源網格，第二個網格是源網格的變形版本。在粗略的偽代碼中，計算過程如下：迭代源網格中的所有頂點。對于每個頂點，計算其與其相鄰頂點的協方差矩陣。對計算出的協方差矩陣執行 SVD 并找到頂點的旋轉矩陣。使用計算出的旋轉矩陣、原始網格中的點坐標和變形網格中的相應坐標來計算頂點的能量偏差。因此，這個能量函數要求我迭代網格中的每個點，并且網格可能有超過 2k 個這樣的點。在 Tensorflow 中，有兩種方法可以做到這一點。我可以有 2 個形狀 (N,3) 的張量，一個代表源點，另一個代表變形網格。純粹使用 Tensorflow 張量來做。也就是說，tf.gather使用現有的 TF 操作迭代上述張量的元素并在每個點上執行計算。這種方法，會非常慢。我之前曾嘗試定義迭代超過 1000 個點的損失函數，而圖構建本身需要太多時間而無法實用。按照此處的 TF 文檔中的說明添加新的 TF OP 。這涉及在 CPP（和 Cuda，用于 GPU 支持）中編寫函數，并使用 TF 注冊新 OP。第一種方法很容易編寫，但速度很慢。第二種方法寫起來很痛苦。我已經使用 TF 3 年了，之前從未使用過 PyTorch，但此時我正在考慮切換到它，如果它為這種情況提供了更好的替代方案。PyTorch 是否有一種方法可以輕松實現此類損失函數，并且執行速度與在 GPU 上一樣快。即，編寫我自己的在 GPU 上運行的損失函數的 pythonic 方式，我沒有任何 C 或 Cuda 代碼？