1. #!/usr/bin/env python
    2. from __future__ import print_function
    3. from itertools import count
    5. import torch
    6. import torch.nn.functional as F
    8. POLY_DEGREE = 4
    9. W_target = torch.randn(POLY_DEGREE, 1) * 5
    10. b_target = torch.randn(1) * 5
    13. def make_features(x):
    14.     """Builds features i.e. a matrix with columns [x, x^2, x^3, x^4]."""
    15.     x = x.unsqueeze(1)
    16.     return torch.cat([x ** i for i in range(1, POLY_DEGREE+1)], 1)
    19. def f(x):
    20.     """Approximated function."""
    21.     return x.mm(W_target) + b_target.item()
    24. def poly_desc(W, b):
    25.     """Creates a string description of a polynomial."""
    26.     result = 'y = '
    27.     for i, w in enumerate(W):
    28.         result += '{:+.2f} x^{} '.format(w, i + 1)
    29.     result += '{:+.2f}'.format(b[0])
    30.     return result
    33. def get_batch(batch_size=32):
    34.     """Builds a batch i.e. (x, f(x)) pair."""
    35.     random = torch.randn(batch_size)
    36.     x = make_features(random)
    37.     y = f(x)
    38.     return x, y
    41. # Define model
    42. fc = torch.nn.Linear(W_target.size(0), 1)
    44. for batch_idx in count(1):
    45.     # Get data
    46.     batch_x, batch_y = get_batch()
    48.     # Reset gradients
    49.     fc.zero_grad()
    51.     # Forward pass
    52.     output = F.smooth_l1_loss(fc(batch_x), batch_y)
    53.     loss = output.item()
    55.     # Backward pass
    56.     output.backward()
    58.     # Apply gradients
    59.     for param in fc.parameters():
    60.         param.data.add_(-0.1 * param.grad)
    62.     # Stop criterion
    63.     if loss < 1e-3:
    64.         break
    66. print('Loss: {:.6f} after {} batches'.format(loss, batch_idx))
    67. print('==> Learned function:\t' + poly_desc(fc.weight.view(-1), fc.bias))
    68. print('==> Actual function:\t' + poly_desc(W_target.view(-1), b_target))