Алгоритм стохастического градиентного спуска с использованием pyspark для получения оптимальных параметров

Question

Я хочу реализовать алгоритм в pyspark, чтобы получить функцию стоимости / значение ошибки, используя как стохастический градиентный спуск, так и градиентный спуск. Я реализовал следующий код в Python, используя модуль Numpy:

from numpy import *
from numpy.random import *
import sys

# Gradient descent function
def GDFunc(weight): 
    return sum((x.dot(weight) - y)**2 for x, y in data_points) / len(data_points)

def dGDFunc(weight):
    return sum(2 * (x.dot(weight) - y) * x for x, y in data_points) / len(data_points)

# Stochastic gradient descent function
def SGDFunc(weight, i):
    x, y = data_points[i]
    return (x.dot(weight) - y)**2

def dSGDFunc(weight, i):
    x, y = data_points[i]
    return 2 * (x.dot(weight) - y) * x


# Data generation
true_weights = array([1, 2, 3, 4, 5])
d = len(true_weights)
data_points = []
for _ in range(1000):
    x = randn(d)
    y = x.dot(true_weights) + randn()
    data_points.append((x, y))

# Find gradient descent optimal parameters
def gradientDescent(GDFunc, dGDFunc, d):
    weight = zeros(d)
    alpha = 0.01
    for t in range(100):
        value = GDFunc(weight)
        gradient = dGDFunc(weight)
        print 'iteration %d: weight = %s, GDFunc(weight) = %s' % (t, weight, value)
        weight = weight - alpha * gradient

def stochasticGD(SGDFunc, dSGDFunc, d, n):
    weight = zeros(d)
    numUpdates = 0
    oldValue=0
    for t in range(100):
        for i in range(n):  # Nested Loop to process each sample individually
            value = SGDFunc(weight, i)
            gradient = dSGDFunc(weight, i)
            numUpdates += 1
            alpha = 1.00 / numUpdates  # Decreasing step size
            weight = weight - alpha * gradient
        print 'iteration %d: Cost: %s.  weight = %s' % (t,value, weight)
        if t>5 and value>oldValue:
            print "Found Local Minima"
            sys.exit()
        oldValue=value


gradientDescent(GDFunc, dGDFunc, d)
stochasticGD(SGDFunc, dSGDFunc, d, len(data_points))

Если я увеличу количество итераций в приведенном выше коде, выполнение займет много времени, поэтому я ищу способ реализовать то же самое с помощью pyspark.

Хотя мне известна библиотека pllpark mllib, которая предоставляет класс LinearRegressionWithSGD для обучения модели, мне неясно, как из нее получить оптимальные параметры.