用PyTorch实现线性回归

用PyTorch实现线性回归

一个小注意点：线性模型中w一般放在矩阵x右边。

即x(n * 3) * w(3 * 2) = y(n * 2)

话不多说，直接上代码：

import torch

# 初始化数据集
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])

# 定义模型
class LinearModel(torch.nn.Module):
    # 需要定义__init__和forward两个函数
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1) # 分别表示x和y的size

    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred
model = LinearModel()

# 损失函数
criterion = torch.nn.MSELoss(size_average=False) # 这个参数是是否求平均
# 优化器，用于不断优化参数w
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # lr是学习率

# 训练
epoch_list = [] #记录轮数
cost_list = [] #记录损失
for epoch in range(200):
    # 第一步，求y_hat
    y_pred = model(x_data)
    # 第二步，求损失loss
    loss = criterion(y_pred, y_data)
    print(epoch, loss)
    epoch_list.append(epoch)
    cost_list.append(loss.item())
    # 第2.5步，梯度清零
    optimizer.zero_grad()
    # 第三步，反向传播
    loss.backward()
    # 第四步，更新参数
    optimizer.step()

print('w = ', model.linear.weight.item())
print('b = ', model.linear.bias.item())

# 测试
x_test = torch.Tensor([[4.0]])
y_test = model(x_test)
print('y_pred = ', y_test.data)

# 画图
import matplotlib.pyplot as plt
plt.plot(epoch_list, cost_list)
plt.ylabel('cost')
plt.xlabel('epoch')
plt.show()

Python函数用法的小提示:

1、参数*args和**kwargs

*args表示多个变量，比如func(1, 2, 3, 4…)

**kwargs表示讲变量抽象成字典，比如func(x = 1, y = 2)

2、函数的使用

# balabala,,,略
def function name (parameter):
   function body
   return