线性回归demo

一个使用pytorch框架训练线性回归的小demo

• 用cpu就能很快跑出来
• 不需要额外的输入文件

import torch
import numpy as np
import torch.nn as nnx_values = [i for i in range(11)] # [0,1,2,3,4,5,6,7,8,9,10]
x_train = np.array(x_values, dtype=np.float32)
x_train = x_train.reshape(-1, 1) # 将x_train调整为11*1的矩阵y_values = [2.5*i+3.5 for i in x_values] # y=2.5x+3.5
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)class LinearRegressionModel(nn.Module): # 继承自nn包的Module类def __init__(self, input_dim, output_dim):super(LinearRegressionModel, self).__init__() # 执行父类的构造函数self.linear = nn.Linear(input_dim, output_dim)# nn.Linear(输入数据维度, 输出数据维度) 全连接层def forward(self, x):out = self.linear(x)return outinput_dim = 1
output_dim = 1
model = LinearRegressionModel(input_dim, output_dim)epochs = 1000 # 训练次数
learning_rate = 0.01 # 学习率
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
# 优化器，使用基本的优化器SGD，传入需要更新的参数（model中的全部参数）和学习率
criterion = nn.MSELoss() # 回归任务可选用MSE等for epoch in range(epochs):epoch += 1# x_train和y_train均为numpy.ndarry格式，需要转换为tensor格式才可以传入框架inputs = torch.from_numpy(x_train)labels = torch.from_numpy(y_train)# 每次迭代开始时 梯度需要清零optimizer.zero_grad()# 前向传播outputs = model.forward(inputs)# 计算损失loss = criterion(outputs, labels)# 反向传播loss.backward()# 更新权重参数optimizer.step()# 每50个epoch输出一次，以显示训练进度if epoch % 50 == 0:print('epoch {}, loss {}'.format(epoch, loss.item()))y_predicted = model.forward(torch.from_numpy(x_train)).data.numpy()
# 前向传播 传入训练数据x 输出预测结果y 用以测试
# .data.numpy() 将结果转换成numpy
print(y_predicted)