线性回归是机器学习中的基础算法之一，通过最小化预测值和真实值之间的误差来拟合数据。在本教程中，我们将使用 PyTorch 从零开始构建一个线性回归模型，并拆解成关键部分，以便更好地理解 PyTorch 的使用方式。

1. 构建组件

在实现线性回归之前，我们需要构建一些核心组件，包括决策函数、损失函数和优化方法。

1.1 构建数据

我们使用 make_regression 生成一个带有噪声的一维回归数据集，并将其划分为训练集和测试集。

import pickle
import numpy as np
from sklearn.datasets import make_regression
import torch


def create_dataset():
    X, y = make_regression(n_samples=256,
                           n_features=1,
                           noise=10,
                           random_state=0)

    # 打乱数据
    indices = np.arange(X.shape[0])  # 获取索引
    np.random.shuffle(indices)  # 打乱索引

    X_shuffled = X[indices]  # 打乱后的特征
    y_shuffled = y[indices]  # 打乱后的标签

    # 数据集分割
    train_size = 0.8
    train_number = int(X_shuffled.shape[0] * 0.8)
    X_train = X_shuffled[:train_number]
    y_train = y_shuffled[:train_number]

    X_test = X_shuffled[train_number:]
    y_test = y_shuffled[train_number:]

    # 转换为 PyTorch 张量
    X_train = torch.tensor(X_train, dtype=torch.float64)
    y_train = torch.tensor(y_train, dtype=torch.float64)

    X_test = torch.tensor(X_test, dtype=torch.float64)
    y_test = torch.tensor(y_test, dtype=torch.float64)

    # 存储数据
    pickle.dump({'data': X_train, 'target': y_train}, open('train.pkl', 'wb'))
    pickle.dump({'data': X_test, 'target': y_test}, open('test.pkl', 'wb'))


if __name__ == '__main__':
    create_dataset()

1.2 数据加载

由于 PyTorch 常用小批量梯度下降（SGD），我们实现一个数据加载器来获取小批量数据。

class DataLoader:

    def __init__(self, X, y, batch_size):
        self.X = X
        self.y = y
        self.batch_size = batch_size

    def __call__(self):
        data_len = len(self.y)
        data_index = list(range(data_len))
        random.shuffle(data_index)
        batch_number = data_len // batch_size

        for idx in range(batch_number):
            start = idx * batch_size
            end = start + batch_size
            yield self.X[start:end], self.y[start:end]

1.3 算法模型

线性回归是 y = wx + b，其中 w 和 b 是需要学习的参数。

class LinearRegression:

    def __init__(self):
        self.w = torch.tensor(0.1, requires_grad=True, dtype=torch.float64)
        self.b = torch.tensor(0.0, requires_grad=True, dtype=torch.float64)

    def __call__(self, x):
        return self.w * x + self.b

    def get_parameters(self):
        return {'w': self.w, 'b': self.b}

    def eval(self):
        self.w.requires_grad = False
        self.b.requires_grad = False

1.4 损失函数

我们使用均方误差（MSE）来衡量预测值和真实值之间的差距。

class SquareLoss:
    def __call__(self, y_pred, y_true):
        return torch.mean((y_pred - y_true) ** 2)

1.5 优化方法

优化器负责参数梯度的清零、以及参数更新，我们这里使用随机梯度下降来更新参数。

class SGD:
    def __init__(self, parameters, lr=1e-2):
        self.lr = lr
        self.parameters = parameters

    def step(self):
        for name, param in self.parameters.items():
            param.data = param.data - self.lr * param.grad.data / 16

    def zero_grad(self):
        for name, param in self.parameters.items():
            param.data.zero_()

2. 训练评估

有了上述组件后，我们就可以训练线性回归模型。

2.1 训练模型

训练过程包括：

• 从数据加载器中获取小批量数据
• 计算预测值
• 计算损失
• 进行反向传播
• 使用 SGD 进行参数更新

import warnings
warnings.filterwarnings('ignore')
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题

from component import LinearRegression
from component import SquareLoss
from component import SGD
from component import data_loader
import pickle


def train():
    # 加载数据
    data = pickle.load(open('train.pkl', 'rb'))
    # 训练参数
    epochs = 100
    estimator = LinearRegression()
    optimizer = SGD(estimator.get_parameters(), lr=1e-2)
    criterion = SquareLoss()

    epoch_loss = []
    for _ in range(epochs):
        total_loss = 0.0
        for train_x, y_true in data_loader(data['data'], data['target'], 16):
            # 前向计算
            y_pred = estimator(train_x)
            # 梯度清零
            optimizer.zero_grad()
            # 损失计算
            loss = criterion(y_pred.reshape(-1), y_true)
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
            # 损失统计
            total_loss += loss.item() * len(y_true)
        epoch_loss.append(total_loss)

    plt.title('损失变化曲线')
    plt.plot(range(len(epoch_loss)), epoch_loss, linestyle='dashed')
    plt.grid()
    plt.show()

    # 存储模型
    pickle.dump(estimator, open('model.pkl', 'wb'))


if __name__ == '__main__':
    train()

2.2 评估模型

在训练完成后，我们使用测试数据集评估模型的表现。

from component import LinearRegression
from component import SquareLoss
from component import SGD
from component import data_loader
import torch
import pickle


def eval():
    # 加载数据
    data = pickle.load(open('test.pkl', 'rb'))
    y_true = data['target']
    estimator = pickle.load(open('model.pkl', 'rb'))
    estimator.eval()
    print(estimator.get_parameters())
    y_pred = estimator(data['data'])
    mse = SquareLoss()(y_pred.reshape(-1), y_true)
    print('MSE:', mse)


if __name__ == '__main__':
    eval()

我们通过手动实现线性回归的假设函数、平方损失、SGD优化方法、以及训练函数来实现对 sklearn make_regression 函数产生的数据集进行拟合，最后通过拟合直线、训练损失变化进行可视化。

import torch
from sklearn.datasets import make_regression
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import random


# 模型参数
w = torch.tensor(0.1, requires_grad=True, dtype=torch.float64)
b = torch.tensor(0.0, requires_grad=True, dtype=torch.float64)


# 1. 数据集函数
def create_dataset():

    x, y, coef = make_regression(n_samples=100,
                                 n_features=1,
                                 noise=10,
                                 coef=True,
                                 bias=14.5,
                                 random_state=0)

    # 转换为张量
    x = torch.tensor(x)
    y = torch.tensor(y)

    return x, y, coef


# 2. 构建数据加载器
def data_loader(x, y, batch_size):

    data_len = len(y)
    data_index = list(range(data_len))
    random.shuffle(data_index)
    batch_number = data_len // batch_size

    for idx in range(batch_number):

        start = idx * batch_size
        end = start + batch_size

        batch_train_x = x[start: end]
        batch_train_y = y[start: end]

        yield batch_train_x, batch_train_y


# 3. 假设函数
def linear_regression(x):
    return w * x + b


# 4. 损失函数
def square_loss(y_pred, y_true):
    return (y_pred - y_true) ** 2


# 5. 优化方法
def sgd(lr=0.01):
    # 使用批量样本的平均梯度
    w.data = w.data - lr * w.grad.data / 16
    b.data = b.data - lr * b.grad.data / 16


# 6. 训练函数
def train():

    # 加载数据集
    x, y, coef = create_dataset()
    # 定义训练参数
    epochs = 100
    learning_rate = 0.01
    # 存储损失
    epoch_loss = []
    total_loss = 0.0
    train_sample = 0

    for _ in range(epochs):

        for train_x, train_y in data_loader(x, y, 16):

            # 训练数据送入模型
            y_pred = linear_regression(train_x)

            # 计算损失值
            loss = square_loss(y_pred, train_y.reshape(-1, 1)).sum()
            total_loss += loss.item()
            train_sample += len(train_y)

            # 梯度清零
            if w.grad is not None:
                w.grad.data.zero_()

            if b.grad is not None:
                b.grad.data.zero_()

            # 反向传播
            loss.backward()

            # 更新参数
            sgd(learning_rate)

            print('loss: %.10f' % (total_loss / train_sample))

        epoch_loss.append(total_loss / train_sample)


    # 绘制拟合直线
    print(coef, w.data.item())
    plt.scatter(x, y)

    x = torch.linspace(x.min(), x.max(), 1000)
    y1 = torch.tensor([v * w + 14.5 for v in x])
    y2 = torch.tensor([v * coef + 14.5 for v in x])

    plt.plot(x, y1, label='训练')
    plt.plot(x, y2, label='真实')
    plt.grid()
    plt.legend()
    plt.show()

    # 打印损失变化曲线
    plt.plot(range(epochs), epoch_loss)
    plt.title('损失变化曲线')
    plt.grid()
    plt.show()


if __name__ == '__main__':
    train()