三默网为您带来有关“pytorch实现线性回归”的文章内容,供您阅读参考。
pytorch实现线性回归
2023-01-21 19:30:36
coding=utf-8
“”"
author:lei
function:
“”"
from torch.optim import SGD
import torch
from torch import nn
定义一个device对象 如果要转换为gpu上运行
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
准备数据
x = torch.rand([500, 1]).to(device)
y_true = 3 * x + 0.8
1、定义模型
class MyLinear(nn.Module):
def init(self):
# 继承父类的init
super(MyLinear, self).init()
self.linear = nn.Linear(1, 1)
def forward(self, x):
out = self.linear(x)
return out
2、实例化模型,优化器实例化,loss实例化 将模型进行gpu实例化
my_linear = MyLinear().to(device)
optimizer = SGD(my_linear.parameters(), 0.001)
loss_fn = nn.MSELoss()
3、循环,进行梯度下降,参数更新
for i in range(120000):
# 得到预测值
y_predict = my_linear(x)
loss = loss_fn(y_predict, y_true)
# 梯度置为一
optimizer.zero_grad()
# 反向传播
loss.backward()
# 参数更新
optimizer.step()
if i % 500 == 0:
params = list(my_linear.parameters())
print(loss.item(), params[0].item(), params[1].item())
4、模型评估 设置模型评估模式,即预测模式
my_linear.eval()