Pytorch入门实例

数据集是受教育年限和收入,如下图

代码如下

import torch
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import torch.nn as nn
import torch.optim as optim
 
data = pd.read_csv('./Income.csv')
 
X = torch.from_numpy(data.Education.values.reshape(-1,1).astype(np.float32))
Y = torch.from_numpy(data.Income.values.reshape(-1,1).astype(np.float32))
 
learning_rate = 0.0001
model = nn.Linear(1,1)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(),lr=learning_rate)
 
for epoch in range(50):
    for x,y in zip(X,Y):
        output = model(x)
        loss = loss_fn(output,y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
 
plt.scatter(data.Education,data.Income)
plt.plot(X.numpy(),model(X).detach().numpy(),c='r')
plt.xlabel('Education')
plt.ylabel('Income')
plt.show()

输出如下