线性回归-pytorch代码

import numpy as np
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F


class LinearRegression(nn.Module):
    def __init__(self, input_size, output_size):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(input_size, output_size)

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


def build_data():
    x = np.array([3, 4, 5, 6, 7, 8, 9], dtype=np.float32)
    x = x.reshape(-1, 1)
    x = Variable(torch.from_numpy(x))

    y = np.array([7.5, 7, 6.5, 6.0, 5.5, 5.0, 4.5], dtype=np.float32)
    y = np.reshape(y, (-1, 1))
    y = Variable(torch.from_numpy(y))
    return x, y


def run_train():
    fit_params = {
        "n_epochs": 100,
    }

    x, y = build_data()
    model = LinearRegression(input_size=1, output_size=1)

    mse = nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.02)

    for i in range(fit_params["n_epochs"]):
        optimizer.zero_grad()
        results = model(x)
        loss = mse(results, y)

        loss.backward()
        optimizer.step()