使用pytorch实现一个线性回归训练函数

使用sklearn.dataset 的make_regression创建用于线性回归的数据集

python 复制代码
def create_dataset():
    x, y, coef = make_regression(n_samples=100, noise=10, coef=True, bias=14.5, n_features=1, random_state=0)
    return torch.tensor(x), torch.tensor(y), coef

加载数据集,并拆分batchs训练集

python 复制代码
def load_dataset(x, y, batch_size):
    data_len = len(y)
    batch_num = data_len // batch_size
    for idx in range(batch_num):
        start = idx * batch_num
        end = idx * batch_num + batch_num
        train_x = x[start : end]
        train_y = y[start : end]
        yield train_x, train_y

定义初始权重和定义计算函数

python 复制代码
w = torch.tensor(0.1, requires_grad=True, dtype=torch.float64)
b = torch.tensor(0, requires_grad=True, dtype=torch.float64)
def linear_regression(x):
    return x * w + b

损失函数使用平方差

python 复制代码
def linear_loss(y_pred, y_true):
    return (y_pred - y_true) ** 2

优化参数使用梯度下降方法

python 复制代码
def sgd(linear_rate, batch_size):
    w.data = w.data - linear_rate * w.grad / batch_size
    b.data = b.data - linear_rate * b.grad / batch_size

训练代码

python 复制代码
def train():
    # 加载数据
    x, y, coef = create_dataset()
    data_len = len(y)

    # 定义参数
    batch_size = 10
    epochs = 100
    linear_rate = 0.01

    # 记录损失值
    epochs_loss = []

    # 迭代
    for eid in range(epochs):
        total_loss = 0.0
        for train_x, train_y in load_dataset(x, y, batch_size):
            # 输入模型
            y_pred = linear_regression(train_x)

            # 计算损失
            loss_num = linear_loss(y_pred, train_y.reshape(-1,1)).sum()

            # 梯度清理
            if w.grad is not None:
                w.grad.zero_()
            if b.grad is not None:
                b.grad.zero_()

            # 反向传播
            loss_num.backward()

            # 更新权重
            sgd(linear_rate, batch_size)

            # 统计损失数值
            total_loss = total_loss + loss_num.item()

        # 记录本次迭代的平均损失
        b_loss = total_loss / data_len
        epochs_loss.append(b_loss)
        print("epoch={},b_loss={}".format(eid, b_loss))

    # 显示预测线核真实线的拟合关系
    print(w, b)
    print(coef, 14.5)

    plt.scatter(x, y)

    test_x = torch.linspace(x.min(), x.max(), 1000)
    y1 = torch.tensor([v * w + b for v in test_x])
    y2 = torch.tensor([v * coef + 14.5 for v in test_x])
    plt.plot(test_x, y1, label='train')
    plt.plot(test_x, y2, label='true')
    plt.grid()
    plt.show()

    # 显示损失值变化曲线
    plt.plot(range(epochs), epochs_loss)
    plt.show()

拟合显示还不错

损失值在低5次迭代后基本就很小了

相关推荐
CoovallyAIHub2 小时前
中科大DSAI Lab团队多篇论文入选ICCV 2025,推动三维视觉与泛化感知技术突破
深度学习·算法·计算机视觉
CoovallyAIHub2 小时前
开源的消逝与新生:从 TensorFlow 的落幕到开源生态的蜕变
pytorch·深度学习·llm
CoovallyAIHub8 小时前
港大&字节重磅发布DanceGRPO:突破视觉生成RLHF瓶颈,多项任务性能提升超180%!
深度学习·算法·计算机视觉
CoovallyAIHub8 小时前
英伟达ViPE重磅发布!解决3D感知难题,SLAM+深度学习完美融合(附带数据集下载地址)
深度学习·算法·计算机视觉
惯导马工1 天前
【论文导读】ORB-SLAM3:An Accurate Open-Source Library for Visual, Visual-Inertial and
深度学习·算法
隐语SecretFlow2 天前
国人自研开源隐私计算框架SecretFlow,深度拆解框架及使用【开发者必看】
深度学习
Billy_Zuo2 天前
人工智能深度学习——卷积神经网络(CNN)
人工智能·深度学习·cnn
羊羊小栈2 天前
基于「YOLO目标检测 + 多模态AI分析」的遥感影像目标检测分析系统(vue+flask+数据集+模型训练)
人工智能·深度学习·yolo·目标检测·毕业设计·大作业
l12345sy2 天前
Day24_【深度学习—广播机制】
人工智能·pytorch·深度学习·广播机制
九章云极AladdinEdu3 天前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力