通过调用Pytorch Api 实现线性回归

最新推荐文章于 2024-08-01 21:08:58 发布

悠悠小船长

最新推荐文章于 2024-08-01 21:08:58 发布

阅读量81

点赞数

分类专栏： DeepLearning 文章标签： pytorch 线性回归深度学习

本文链接：https://blog.csdn.net/WakingStone/article/details/126906735

版权

DeepLearning 专栏收录该内容

19 篇文章 1 订阅

订阅专栏

具体代码

import torch
import torch.nn as nn
from torch.optim import SGD

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.rand([500, 1]).to(device)
y_true = 6 * x + 2.5


class MyLinear(nn.Module):
    def __init__(self):
        super(MyLinear, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        out = self.linear(x)
        return out


model = MyLinear().to(device)
optimizer = SGD(model.parameters(), 0.001)
loss_fn = nn.MSELoss()

for epoch in range(50000):
    y_predict = model(x)
    loss = loss_fn(y_predict, y_true)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 100 == 0:
        params = list(model.parameters())
        print(f'loss:{loss.item()} {params[0].item()} {params[1].item()}')