我们通过手动实现线性回归的假设函数、平方损失、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()
输出结果:
loss: 1296.1368728541 loss: 2436.9092349261 loss: 2223.9904977827 loss: 2226.8678712042 loss: 2346.9798105877 loss: 2059.6439491472 loss: 1913.1633767652 loss: 2027.7041992174 loss: 1958.8407893180 loss: 1936.7456442893 loss: 1958.9884082751 loss: 1836.6773229252 ... loss: 200.2883544967 loss: 200.1779015367 loss: 200.0914139110 loss: 199.9625017734 loss: 199.8119659842 loss: 199.5681193811 loss: 199.4116960121 loss: 199.3038260199 loss: 199.2196718821 loss: 199.0935106288 loss: 198.9459334305 loss: 198.7059688806 42.38550485581797 42.887112362029214 Process finished with exit code 0