pytorch实现线性回归
导入相关python包
import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from torch import nn
%matplotlib inline
加载数据
data = pd.read_csv('E:/datasets/dataset/Income1.csv')
X = torch.from_numpy(data.Education.values.reshape(-1, 1).astype(np.float32)) #DataFrame转tensor的常用方法
Y = torch.from_numpy(data.Income.values.reshape(-1, 1).astype(np.float32))
![](https://img2022.cnblogs.com/blog/2281505/202202/2281505-20220207094222264-1004561779.png)
定义模型
model = nn.Linear(in_features = 1, out_features = 1) # w * input + b 等价于 model(input)
loss_func = nn.MSELoss() # 损失函数
optimizer = torch.optim.SGD(params = model.parameters(), lr = 0.0001)
训练模型
for epoch in range(5000):
for x, y in zip(X, Y):
y_pred = model(x) # 使用模型预测
loss = loss_func(y, y_pred) # 根据预测结果计算损失
optimizer.zero_grad() # 把变量梯度清 0
loss.backward() # 求解梯度
optimizer.step() # 优化模型参数
查看训练结果
model.weight, model.bias
![](https://img2020.cnblogs.com/blog/2281505/202101/2281505-20210126105635697-1938725774.png)
plt.scatter(data.Education, data.Income)
plt.plot(X.numpy(), model(X).data.numpy(), c = 'r')
#model(X).data是将tensor取出,否则会带着grad和grad_fn
![](https://img2020.cnblogs.com/blog/2281505/202101/2281505-20210126105652027-1736637314.png)