问基于Pandas + GIF可视化的梯度下降算法

EN

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import random
import imageio
import os


def compute_cost(b, m, data):
    """
    Compute cost function for univariate linear regression using mean squared error(MSE).
    Args:
        b: Intercept (y = mx + b).
        m: Slope (y = mx + b).
        data: A pandas df with x and y data.
    Return:
         Cost function value.
    """
    data['Mean Squared Error(MSE)'] = (data['Y'] - (m * data['X'] + b)) ** 2
    return data['Mean Squared Error(MSE)'].sum().mean()


def adjust_gradient(b_current, m_current, data, learning_rate):
    """
    Adjust Theta parameters for the univariate linear equation y^ = hθ(x) = θ0 + θ1x
    Args:
        b_current: Current Intercept (y = mx + b) or [Theta(0) θ0].
        m_current: Current Slope (y = mx + b) or [Theta(1) θ1].
        data: A pandas df with x and y data.
        learning_rate: Alpha value.
    Return:
        Adjusted Theta parameters.
    """
    data['b Gradient'] = -(2 / len(data)) * (data['Y'] - ((m_current * data['X']) + b_current))
    data['m Gradient'] = -(2 / len(data)) * data['X'] * (data['Y'] - ((m_current * data['X']) + b_current))
    new_b = b_current - (data['b Gradient'].sum() * learning_rate)
    new_m = m_current - (data['m Gradient'].sum() * learning_rate)
    return new_b, new_m


def gradient_descent(data, b, m, learning_rate, max_iter, visual=False):
    """
    Optimize Theta values for the univariate linear regression equation y^ = hθ(x) = θ0 + θ1x.
    Args:
        data: A pandas df with x and y data.
        b: Starting b (θ0) value.
        m: Starting m (θ1) value.
        learning_rate: Alpha value.
        max_iter: Maximum number of iterations.
        visual: If True, a GIF progression will be generated.
    Return:
        Optimized values for θ0 and θ1.
    """
    line = np.arange(len(data))
    folder_name = None
    if visual:
        folder_name = str(random.randint(10 ** 6, 10 ** 8))
        os.mkdir(folder_name)
        os.chdir(folder_name)
    for i in range(max_iter):
        b, m = adjust_gradient(b, m, data, learning_rate)
        if visual:
            data['Line'] = (line * m) + b
            data.plot(kind='scatter', x='X', y='Y', figsize=(8, 8), marker='x', color='r')
            plt.plot(data['Line'], color='b')
            plt.grid()
            plt.title(f'y = {m}x + {b}\nCurrent cost: {compute_cost(b, m, data)}\nIteration: {i}\n'
                      f'Alpha = {learning_rate}')
            fig_name = ''.join([str(i), '.png'])
            plt.savefig(fig_name)
            plt.close()
    if visual:
        frames = os.listdir('.')
        frames.sort(key=lambda x: int(x.split('.')[0]))
        frames = [imageio.imread(frame) for frame in frames]
        imageio.mimsave(folder_name + '.gif', frames)
    return b, m


if __name__ == '__main__':
    data = pd.read_csv('data.csv')
    data.columns = ['X', 'Y']
    learning = 0.00001
    initial_b, initial_m = 0, 0
    max_it = 350
    b, m = gradient_descent(data, initial_b, initial_m, learning, max_it, visual=True)