三默网为您带来有关“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()