pytorchGPU|CPU自动推断版线性回归demo

from sklearn.datasets import load_diabetes
from sklearn.utils import shuffle
import torch
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np

def is_GPU(tensor):
    '''
    函数名 : cpu 或者 gpu 自动推断函数
    :param tensor: tensor结构数据
    :return: 根据系统转cpu或者gpu结构
    '''
    if torch.cuda.is_available():
        return tensor.cuda()
    return tensor

def yhat_function(X,theta,b): # 期望函数
    '''
    函数名 : yhot函数 / 模型的期望的y值的函数
    :param X: 数据集
    :param a: 待迭代参数
    :param b: 偏置选项
    :return: 模型的线性期望
    '''
    return is_GPU(X.mm(theta)+b)

def initialize_theta(dims,b=None): # 初始化参数
    '''
    函数名 : 初始化参数函数
    :param dims: 数据集的列数
    :return: thate -> 初始化的参数 ， b -> 初始化偏置
    '''
    
    theta = is_GPU(torch.zeros((dims, 1)))
    if b == None :
        b = 0
    return theta, b


def loss_function(y_hat,y,m): # 损失函数
    '''
    函数名 : 损失函数
    :param y_hat: 模型当前拟合的y值
    :param y: 真实的y值
    :param m: 数据集的行数
    :return: 损失值
    '''
    return is_GPU(((y_hat-y)**2).sum()/m)

def partial_derivative_function(X,yhat,Y,m): # 求偏导数
    '''
    函数名 : 求偏导数函数
    :param X: 数据集
    :param yhat: y值的拟合过程
    :param Y: 真实的y值
    :param m: 数据集X的行数
    :return: 偏导数
    '''
    dw = X.T.mm(yhat-Y)/m
    db = (yhat-Y).sum()/m
    return dw,db

def liner_model(X_train,y_train,b=None
                ,maxloop = 10000
                ,alpha = 0.001
                ,convergence_threshold = 10**(-8)
                ,display=False,display_cycle=False):
    '''
    :param X_train: 训练集X
    :param y_train: 训练标签 y
    :param b:偏置项
    :param maxloop: 最大迭代次数
    :param alpha: 学习率
    :param convergence_threshold: 收敛阈值
    :param display:是否显示迭代过程与损失函数图像
    :param display_cycle: 显示迭代过程的间隔，默认每1/10次显示
    :return: lossdata -> 损失函数过程收集,loss ->最终损失值 , params —> 参数 , grads -> 梯度值 ,gradsdata -> 梯度迭代函数
    '''

    theta,b = initialize_theta(X_train.shape[1],b)
    lossdata = [np.Inf]
    m,n = X_train.shape
    gradsdata=[]
    for i in range(1,maxloop):
        yhat = yhat_function(X_train,theta,b)
        loss = loss_function(yhat,y_train,m)
        dw,db = partial_derivative_function(X_train,yhat,y_train,m)
        # 参数更新过程
        theta += is_GPU(-alpha * dw)
        b += is_GPU(-alpha * db)
        # 保存参数
        params = {'theta': theta,'b': b}
        # 保存梯度
        grads = {'dw': dw,'db': db}
        
        gradsdata.append([dw,db])
        lossdata.append(loss)
        if display:
            cycle = display_cycle == False and maxloop/10 or display
            if i % cycle == 0:
                print('epoch %d loss %f' % (i, loss))
        if lossdata[-2]-lossdata[-1] <= convergence_threshold:
            break
    if display:
        plt.plot(lossdata)
        plt.show()
        plt.close()
    print(i)
    return lossdata,loss, params, grads,gradsdata

def predict(X, params):
    '''
    :param X:
    :param params:
    :return:
    '''
    w = params['theta']
    b = params['b']
    y_pred = X.mm(w) + b
    return y_pred