Используйте Tensorflow/PyTorch для ускорения минимизации пользовательской функции.

Я делаю много симуляций, для которых мне часто нужно минимизировать сложные пользовательские функции, для которых я обычно использую numpy и scipy.optimize.minimize(). Однако проблема в том, что мне нужно явно записать функцию градиента, которую иногда бывает очень сложно/невозможно найти. А для векторов большой размерности числовые производные, вычисляемые scipy, непомерно дороги.

Итак, я пытаюсь переключиться на Tensorflow или PyTorch, чтобы воспользоваться как их возможностями автоматической дифференциации, так и иметь возможность свободно использовать графические процессоры. Позвольте мне привести явный пример функции, производную которой довольно сложно записать (потребуется много цепных правил), и, таким образом, кажется, что она созрела для 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.

python pytorch tensorflow optimization minimization
person ap21    schedule 06.07.2019    source источник
comment
1) вам, вероятно, понадобится что-то вроде пакетов jax/autograd. 2) у меня тоже нет   -  person Shihab Shahriar Khan    schedule 08.07.2019

Ответы (2)


arrow_upward
1
arrow_downward

Вот решение, использующее автоматическое вычисление градиента с помощью факела, а затем минимизатор scipy с библиотекой, которую я написал autograd -свернуть. Преимуществом перед SGD является более высокая точность оценки (при использовании методов второго порядка). Вероятно, это эквивалентно использованию 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]))
person bruno    schedule 08.04.2021

arrow_upward
0
arrow_downward

Я работаю над тем же. Вот что я получил.

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) > 0.001, -angle, angle)
    angle = torch.where(compute_angle(cb, w) > 1, -angle, angle)
#     try:
#         if compute_angle(cb, w) > 0.001:
#             angle = -angle
#     except ZeroDivisionError:
#         # dihedral=pi
#         pass
    return angle

v1 = torch.tensor([-17.0490,   5.9270,  21.5340], requires_grad=True)
v2 = torch.tensor([-0.1608,  0.0600, -0.0371], requires_grad=True)
v3 = torch.tensor([-0.2000,  0.0007, -0.0927], requires_grad=True)
v4 = torch.tensor([-0.1423,  0.0197, -0.0727], requires_grad=True)

dihedral = compute_dihedral(v1,v2,v3,v4)
target_dihedral = -2

print(dihedral)   # should print -2.6387


for i in range(100):
    dihedral = compute_dihedral(v1,v2,v3,v4)
    loss = (dihedral - target_dihedral)**2
    loss.backward()
    learning_rate = 0.001
    with torch.no_grad():
        v1 -= learning_rate * v1.grad
        v2 -= learning_rate * v2.grad
        v3 -= learning_rate * v3.grad
        v4 -= learning_rate * v4.grad

        # Manually zero the gradients after updating weights
        v1.grad = None
        v2.grad = None
        v3.grad = None
        v4.grad = None
print(compute_dihedral(v1,v2,v3,v4))   # should print -2
person wei    schedule 22.02.2021
comment
Вы когда-нибудь выполняли минимизацию? - person Arya McCarthy; 22.02.2021
comment
Теперь думаю. - person wei; 22.02.2021
comment
Ага. Обратите внимание, что вы можете получить идентичный результат, используя torch.optim.SGD, и вам не придется вручную обрабатывать градиент каждого значения. - person Arya McCarthy; 22.02.2021