import torch
from torch.autograd import Variable
import matplotlib.pyplot as plt
from torch import nnPOLY_DEGREE = 3

def make_features(x):x = x.unsqueeze(1)return torch.cat([x ** i for i in range(1, POLY_DEGREE+1)], 1)

W_target = torch.FloatTensor([3,6,2]).unsqueeze(1)
b_target = torch.FloatTensor([8])def f(x):return x.mm(W_target) + b_target.item()

def get_batch(batch_size=64):random = torch.randn(batch_size)x = make_features(random)y = f(x)  # + torch.rand(1)return Variable(x), Variable(y)# # show an example of a batch
# x_axis = torch.randn(64)
# y_axis = f(make_features(x_axis)).squeeze()
# plt.title("Original Data Example")
# plt.scatter(x_axis.data.numpy(), y_axis.data.numpy())
# plt.show()

class poly_model(nn.Module):def __init__(self):super(poly_model,self).__init__()self.poly = nn.Linear(3,1)def forward(self, x):out = self.poly(x)return outmodel = poly_model()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

def poly_desc(W, b):result = 'y = 'for i, w in enumerate(W):result += '{:+.2f} x^{} '.format(w, len(W) - i)result += '{:+.2f}'.format(b[0])return result

epoch = 0
while True:batch_x, batch_y = get_batch()output = model(batch_x)loss = criterion(output, batch_y)print_loss = loss.dataoptimizer.zero_grad()loss.backward()optimizer.step()# if epoch % 50 == 0:#     print('Loss: {:.6f} after {} batches'.format(loss, epoch))#     print('==> Learned function:t' + poly_desc(model.poly.weight.view(-1), model.poly.bias))epoch += 1if print_loss < 1e-3:print()print("==========End of Training==========")breakprint('Loss: {:.6f} after {} batches'.format(loss, epoch))
print('==> Learned function:t' + poly_desc(model.poly.weight.view(-1), model.poly.bias))
print('==> Actual function:t' + poly_desc(W_target.view(-1), b_target))