用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.linspace(-3, 3, 100000)
X = Produce_X(x)
y = x + 1.2*torch.rand(x.size())
w = torch.rand(2)


plt.scatter(x.numpy(), y.numpy(), s=0.001)
plt.show()

CUDA = torch.cuda.is_available()

if CUDA:
    inputs = X.cuda()
    target = y.cuda()
    w = w.cuda()
    w.requires_grad=True
else:
    inputs = X
    target = y
    w = w
    w.requires_grad = True


def draw(output, loss):
    if CUDA:
        output= output.cpu()
    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()/100000
        loss.backward()
        w.data -= learning_rate * w.grad
        w.grad.zero_()

        if epoch %80 ==0:
            draw(output, loss)

    return w, loss


from time import perf_counter
start = perf_counter()
w, loss = train(10000, learning_rate=1e-4)
finish = perf_counter()
time = finish - start

print('计算时间:%s' % time)
print("final loss:", loss.item())
print("weights:", w.data)