线性回归是机器学习中的基础算法之一,通过最小化预测值和真实值之间的误差来拟合数据。在本教程中,我们将使用 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()



冀公网安备13050302001966号