import numpy as np import torch from torch.autograd import Variable import matplotlib.pyplot as plt """ Using pytorch to do linear regression """ torch.manual_seed(2018) # model's real-parameters w_target = 3 b_target = 10 # generate data n_data = 100 x_train = np.random.rand(n_data, 1)*20 - 10 y_train = w_target*x_train + b_target + (np.random.randn(n_data, 1)*10-5.0) # draw the data plt.plot(x_train, y_train, 'bo') plt.show() # convert to tensor x_train = torch.from_numpy(x_train).float() y_train = torch.from_numpy(y_train).float() # define model parameters w = Variable(torch.randn(1).float(), requires_grad=True) b = Variable(torch.zeros(1).float(), requires_grad=True) # construct the linear model x_train = Variable(x_train) y_train = Variable(y_train) # define model's function def linear_model(x): return x*w + b # define the loss function def get_loss(y_pred, y): return torch.mean((y_pred - y)**2) # upgrade parameters eta = 1e-2 n_epoch = 100 for i in range(n_epoch): y_pred = linear_model(x_train) loss = get_loss(y_pred, y_train) loss.backward() w.data = w.data - eta*w.grad.data b.data = b.data - eta*b.grad.data w.grad.zero_() b.grad.zero_() if i % 10 == 0: print("epoch: %3d, loss: %f" % (i, loss.data[0])) # draw the results plt.plot(x_train.data.numpy(), y_train.data.numpy(), 'bo', label="Real") plt.plot(x_train.data.numpy(), y_pred.data.numpy(), 'ro', label="Estimated") plt.legend() plt.show()