Градиентный спуск для линейной регрессии с mxnet.autograd - проблемы производительности

Question

Я реализовал простой алгоритм градиентного спуска для линейной регрессии, используя mxnet.autograd.
Все работает просто отлично, но производительность ужасна.Я использую ванильный градиентный спуск, а не SGD, но я сомневаюсь, что в этом проблема ... если я просто использую аналитическое выражение для градиента, весь процесс за 1000 эпох занимает около 2 с, но при использовании автограда он поднимается до 147 с.

Это реализация кода

from mxnet import nd, autograd, gluon
import pandas as pd



def main():
    # learning algorithm parameters
    nr_epochs = 1000
    alpha = 0.01

    # read data
    data = pd.read_csv("dataset.txt", header=0, index_col=None, sep="\s+")


    # ---------------------------------
    # --   using gradient descent   ---
    # ---------------------------------
    data.insert(0, "x_0", 1, True)                              # insert column of "1"s as x_0
    m = data.shape[0]                                           # number of samples
    n = data.shape[1] - 1                                       # number of features
    X = nd.array(data.iloc[:, 0:n].values)                      # array with x values
    Y = nd.array(data.iloc[:, -1].values)                       # array with y values

    theta = nd.zeros(n)                                         # initial parameters array
    theta.attach_grad()                                         # declare gradient with respect to theta is needed
    # ----------------------------------------------------------
    theta, Loss = GradientDescent(X, Y, theta, alpha, nr_epochs)
    # ----------------------------------------------------------

    print("Theta by gradient descent:")
    print(theta)


#--------------#
#   END MAIN   #
#--------------#



#-------------------#
#   loss function   #
#-------------------#
def LossFunction(X, Y, theta):
    m = X.shape[0]                  # number of training samples
    loss = 0
    for i in range(X.shape[0]):
        loss = loss + (1 / (2 * m)) * (H(X[i, :], theta) - Y[i]) ** 2
    return loss


#----------------#
#   hypothesis   #
#----------------#
def H(x, theta):
    return nd.dot(x, theta)



#----------------------#
#   gradient descent   #
#----------------------#
def GradientDescent(X, Y, theta, alpha, nr_epochs):

    Loss = nd.zeros(nr_epochs)                                 # array containing values of loss function over iterations

    for epoch in range(nr_epochs):
        with autograd.record():
            loss = LossFunction(X, Y, theta)
        loss.backward()
        Loss[epoch] = loss

        for j in range(len(theta)):
            theta[j] = theta[j] - alpha * theta.grad[j]

    return theta, Loss




if __name__ == "__main__":
    main()

Узким местом является вызов

theta, Loss = GradientDescent(X, Y, theta, alpha, nr_epochs)

Я что-то не так делаю?Я видел несколько других примеров, и они работают намного быстрее, чем мой, есть ли что-то, что я мог бы изменить, чтобы уменьшить время работы?Спасибо!