Улучшить скорость градиентного спуска

Question

Я пытаюсь максимизировать целевую функцию f(x) с функцией scipy.optimize.minimum. Но запуск кода обычно занимает 4-5 часов, потому что функция f(x) требует большого количества вычислений сложной матрицы. Чтобы улучшить его скорость, я хочу использовать gpu. И я уже попробовал tensorflow пакет. Поскольку я использую numpy для определения f(x), я должен преобразовать его в формат тензорного потока. Тем не менее, он не поддерживает вычисление сложной матрицы. Какой еще пакет или средство я могу использовать? Любые предложения?

Чтобы указать c мою проблему, я покажу схему расчета ниже:

1. Рассчитать ожидание :

   - где H=x*H_0, x - параметр

2. Пусть \phi go через динамику уравнения Шредингера

   -Different H соответствует другому \phi_end. Таким образом, параметр x определяет ожидание

3. Изменение x, вычисляет соответствующее ожидание
4. Найдите спецификацию c x, минимизирующую ожидание

Вот простой пример части моего кода:

import numpy as np
import cmath
from scipy.linalg import expm
import scipy.optimize as opt

# create initial complex matrixes
N = 2 # Dimension of matrix
H = np.array([[1.0 + 1.0j] * N] * N) # a complex matrix with shape(N, N)
A = np.array([[0.0j] * N] * N)
A[0][0] = 1.0 + 1j

# calculate the expectation
def value(phi): 
    exp_H = expm(H) # put the matrix in the exp function
    new_phi = np.linalg.linalg.matmul(exp_H, phi)
    # calculate the expectation of the matrix
    x = np.linalg.linalg.matmul(H, new_phi)
    expectation = np.inner(np.conj(phi), x)
    return expectation

# Contants
tmax = 1
dt = 0.1
nstep = int(tmax/dt)
phi_init = [1.0 + 1.0j] * N


# 1st derivative of Schrödinger equation
def dXdt(t, phi, H): # 1st derivative of the function
    return -1j * np.linalg.linalg.matmul(H, phi)

def f(X):
    phi = [[0j] * N] * nstep # store every time's phi
    phi[0] = phi_init
    # phi go through the dynamics of Schrödinger equation
    for i in range(nstep - 1):
        phi[i + 1] =  phi[i] - dXdt(i * dt, X[i] * H, phi[i]) * dt
    # calculate the corresponding value    
    f_result = value(phi[-1])
    return f_result

# Initialize the parameter
X0 = np.array(np.ones(nstep))
results = opt.minimize(f, X0) # minimize the target function
opt_x = results.x

PS: Python Версия: 3.7

Операционная система: Win 10