python - 使用 Tensorflow/PyTorch 加速自定义函数的最小化

Question

我做了很多模拟，为此我经常需要最小化复杂的用户定义函数，为此我通常使用 numpy 和 scipy.optimize.minimize()。然而，这个问题是我需要明确地写下一个梯度函数，这有时很难/不可能找到。对于大维向量，scipy 计算的数值导数非常昂贵。

因此，我正在尝试切换到 Tensorflow 或 PyTorch 以利用它们的自动微分功能，并能够自由利用 GPU。让我给出一个函数的明确示例，该函数的导数写起来有些复杂(需要大量链式法则)，因此对于 Tensorflow 或 PyTorch 来说似乎已经成熟 - - 计算由 3d 空间中的四个点形成的两个三角形之间的二面角:

def dihedralAngle(xyz):
## calculate dihedral angle between 4 nodes

    p1, p2, p3, p4 = 0, 1, 2, 3

    ## get unit normal vectors
    N1 = np.cross(xyz[p1]-xyz[p3] , xyz[p2]-xyz[p3])
    N2 = - np.cross(xyz[p1]-xyz[p4] , xyz[p2]-xyz[p4])
    n1, n2 = N1 / np.linalg.norm(N1), N2 / np.linalg.norm(N2) 

    angle = np.arccos(np.dot(n1, n2))

    return angle

xyz1 = np.array([[0.2       , 0.        , 0.        ],
       [0.198358  , 0.02557543, 0.        ],
       [0.19345897, 0.05073092, 0.        ],
       [0.18538335, 0.0750534 , 0.        ]]) # or any (4,3) ndarray

print(dihedralAngle(xyz1)) >> 3.141

我可以使用 scipy.optimize.minimize() 轻松地最小化它，我应该得到 0。对于这么小的函数，我真的不需要梯度(显式或数值) .但是，如果我希望遍历很多很多节点，并最小化一些依赖于所有二面角的函数，那么开销会高得多吗？

那么我的问题——

1. 如何使用 TensorFlow 或 PyTorch 实现这个最小化问题？对于单个二面角和此类角度的列表(即我们需要考虑列表循环)。
2. 此外，我是否可以使用自动微分获得梯度，并在需要时插入 scipy.optimize.minimize()？例如，scipy.optimize.minimize() 可以轻松实现边界和约束，这是我在 Tensorflow 或 PyToch 优化模块中没有注意到的。

Answer 1

这是一个使用 torch 自动计算梯度的解决方案，然后使用我编写的 lib 的 scipy 最小化器 autograd-minimize .与 SGD 相比的优势是估计精度更高(使用二阶方法)。它可能等同于使用 torch 中的 LBFGS:

import numpy as np
import torch
from autograd_minimize import minimize


def dihedralAngle(xyz):
## calculate dihedral angle between 4 nodes

    p1, p2, p3, p4 = 0, 1, 2, 3

    ## get unit normal vectors
    N1 = np.cross(xyz[p1]-xyz[p3] , xyz[p2]-xyz[p3])
    N2 = - np.cross(xyz[p1]-xyz[p4] , xyz[p2]-xyz[p4])
    n1, n2 = N1 / np.linalg.norm(N1), N2 / np.linalg.norm(N2) 

    angle = np.arccos(np.dot(n1, n2))

    return angle
def compute_angle(p1, p2):
    # inner_product = torch.dot(p1, p2)
    inner_product = (p1*p2).sum(-1)
    p1_norm = torch.linalg.norm(p1, axis=-1)
    p2_norm = torch.linalg.norm(p2, axis=-1)
    cos = inner_product / (p1_norm * p2_norm)
    cos = torch.clamp(cos, -0.99999, 0.99999)
    angle = torch.acos(cos)
    return angle

def compute_dihedral(v1,v2,v3,v4):
    ab = v1 - v2
    cb = v3 - v2
    db = v4 - v3
    u = torch.cross(ab, cb)
    v = torch.cross(db, cb)
    w = torch.cross(u, v)
    angle = compute_angle(u, v)
    angle = torch.where(compute_angle(cb, w) > 1, -angle, angle)

    return angle

def loss_func(v1,v2,v3,v4):
    return ((compute_dihedral(v1,v2,v3,v4)+2)**2).mean()


x0=[np.array([-17.0490,   5.9270,  21.5340]),
    np.array([-0.1608,  0.0600, -0.0371]),
    np.array([-0.2000,  0.0007, -0.0927]),
    np.array([-0.1423,  0.0197, -0.0727])]

res = minimize(loss_func, x0, backend='torch')

print(compute_dihedral(*[torch.tensor(v) for v in res.x]))