import torch
from torch import nn
from torch import optim
import numpy as np
from matplotlib import pyplot as plt
# 1.定义数据
x = torch.randn([50,1])
y = 3*x+0.8
# 2.定义模型
class Lr(nn.Module):
def __init__(self):
super(Lr, self).__init__()
self.linear = nn.Linear(1,1)
def forward(self,x):
out = self.linear(x)
return out
# 3.实例化模型,loss,和优化器
model = Lr()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(),lr=1e-3)
# 4.训练模型
for i in range(3000):
out = model(x) #获取预测值
loss = criterion(y,out) #计算损失
optimizer.zero_grad() #梯度归0
loss.backward() #计算梯度
optimizer.step() #更新梯度
if(i-1)%20 ==0:
print("Epoch[{}/{}],loss:{:.6f}".format(i,3000,loss.data))
# 5.模型评估
model.eval() #设置模型为评估模型,即预测模式
predict = model(x)
predict = predict.data.numpy()
plt.scatter(x.data.numpy(),y.data.numpy(),c="r")
plt.plot(x.data.numpy(),predict)
plt.show()
API实现线性回归
最新推荐文章于 2022-04-15 15:57:07 发布