梯度下降求theta值

import numpy as np
import matplotlib.pyplot as plt

def jfunction(theta0, theta1, x, y, m):
    h = theta0 + theta1 * x
    j, j1 = 0, 0
    for i in range(m):
        j1 += (h[i] - y[i]) ** 2
    j = j1 / (2 * m)
    return j
    
def afunction(theta0, theta1, alph, x, y, m):
    j, j0, j1 = 1, 0, 0
    while abs(j) > 0.198:
        for i in range(m):
            j0 += theta0 + theta1 * x[i] - y[i]
            j1 += (theta0 + theta1 * x[i] - y[i]) * x[i]
        theta0 = theta0 - alph * j0 / m
        theta1 = theta1 - alph * j1 / m
        j = jfunction(theta0, theta1, x, y, m)
        #print(j)
    return theta0, theta1
    
train_set = np.array([[0.5, 2], [1, 1], [1, 2], [1.5, 2.5], [2, 3], [2, 4], [3, 5], [3.5, 6], [4, 5]]).reshape(9, 2)
x_set = train_set[:, 0]#0代表第一列
y_set = train_set[:, -1]#-1代表最后一列
theta0, theta1 = afunction(10, 10, 0.001, x_set, y_set, 9)#afunction(theta0, theta1, alph, x, y, m):
x = np.linspace(0, 5, 100)
print(x)
fx = [theta0 + theta1 * i for i in x]
print(fx)
print("theta0={}\ntheta1={}".format(theta0, theta1))
plt.plot(x_set, y_set, 'ro')
plt.plot(x, fx)
plt.show()