用pytorch实现线性回归

用pytorch实现线性回归

import torch
import matplotlib.pyplot as plt


def Produce_X(x):
    x0 = torch.ones(x.numpy().size)
    X = torch.stack((x, x0), dim=1)
    return X


x = torch.Tensor([1.4, 5, 11, 16, 21])
y = torch.Tensor([14.4, 29.6, 62, 85.5, 113.4])
X = Produce_X(x)
inputs = X
target = y
w = torch.rand(2, requires_grad=True)


def draw(output, loss):
    plt.cla()
    plt.scatter(x.numpy(), y.numpy())
    plt.plot(x.numpy(), output.data.numpy(), 'r-', lw=5)
    plt.text(0.5, 0, 'loss=%s' % (loss.item()), fontdict={'size': 20, 'color': 'red'})
    plt.pause(0.005)


def train(epochs=1, learning_rate=0.01):

    for epoch in range(epochs):
        output = inputs.mv(w)
        loss = (output - target).pow(2).sum()
        loss.backward()
        w.data -= learning_rate * w.grad
        w.grad.zero_()
        if epoch % 80 == 0:
            draw(output, loss)
    return w, loss


w, loss = train(10000, learning_rate=1e-4)
print("final loss:", loss.item())
print("weights:", w.data)