梯度下降法拟合直线

给定300个样本点(x, y)。拟合直线：

设直线方程为y = w*x + b

其中：x，y是给定的样本点，作为训练集。w和b是待拟合的参数。

该问题可以转化为，优化函数 f(w, b) = y - w*x - b

使用平方损失，转化为优化：f(w, b) = (y - w*x -b)^2

求w的梯度为: grad_w = -x * 2(y - w*x -b)

求b的梯度为：grad_b = -1 * 2(y - w*x -b)

再使用：

w = w - lr * grad_w

b = b - lr * grad_b

优化这两个变量。

import numpy as np
import random
import cv2


def fit_line_by_grad():
    gt_x = range(300)
    gt_y = [2*x+4 + 20 * np.random.random() for x in gt_x]
    print("train sample: {}".format(len(gt_x)))

    # y = w * x + b ==> f(w, b) = y - w * x -b ==> f(w, b) = (y - w * x -b)^2
    # ==> grad_w = -x * 2 (y - w * x -b), grad_b = -1 * 2 (y - w * x -b)
    iteration = 10000
    lr = 0.000001
    w = 0.0001
    b = 0
    for i in range(iteration):
        i = i % len(gt_x)
        x = gt_x[i]
        y = gt_y[i]
        grad_w = - x * 2 * (y - w * x -b)
        grad_b = -1 * 2 * (y - w * x -b)

        w = w - lr * grad_w
        b = b - lr * grad_b
    print("y = {}*x + {}".format(w, b))

    img = np.zeros([500, 500, 3], dtype=np.uint8)
    for gt in zip(gt_x, gt_y):
        x, y = [int(i) for i in gt]
        cv2.circle(img, (x, y), 3, (0, 255, 0), 1)
    fit1 = (0, int(w * 0 + b))
    fit2 = (300, int(w * 300 + b))
    cv2.line(img, fit1, fit2, (0, 0, 255), 3)
    cv2.namedWindow("res", 0)
    cv2.imshow("res", img)
    key = cv2.waitKey(0)
    if key == 27:
        exit()


fit_line_by_grad()

最终的拟合效果为：