P5 用Pytorch实现线性回归

75 阅读1分钟
import torch  

#导入数据
x_data = torch.Tensor([[1.0], [2.0], [3.0]])  
y_data = torch.Tensor([[2.0], [4.0], [6.0]])  

#基于Pytorch定义模型
class LinearModel(torch.nn.Module):  
    def __init__(self):  
    super(LinearModel, self).__init__()  
    self.linear = torch.nn.Linear(1, 1)  
  
def forward(self, x):  
    y_pred = self.linear(x)  
    return y_pred  
  
model = LinearModel()  

#构建损失和优化器
criterion = torch.nn.MSELoss(size_average=False)  
optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)  

losses = []
num_epoch = 1000

#模型训练
for epoch in range(num_epoch):  
    y_pred = model(x_data)  
    loss = criterion(y_pred, y_data)  
    losses.append(loss.item())
    print(epoch, loss.item())  
  
    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)

plt.plot(range(num_epoch), losses)
plt.xlabel(Epoch)
plt,ylabel(Losses)
plt,show()

类似的优化算法: torch.optim.Adagrad torch.optim.Adam torch.optim.Adamax torch.optim.ASGD torch.optim.LBFGS torch.optim.RMSprop torch.optim.Rprop torch.optim.SGD

Pytorch官方学习网站:pytorch.org/tutorials/b… 英语不好也可以看看这篇文章: blog.csdn.net/TYtrack/art…