<jupyter_start><jupyter_text><jupyter_code>%matplotlib inline<jupyter_output><empty_output><jupyter_text>Disclaimer: This notebook is modified from the following notebooks:- https://colab.research.google.com/github/pytorch/tutorials/blob/gh-pages/_downloads/nn_tutorial.ipynbscrollTo=N29T2vcqOnVU<jupyter_code>from matplotlib import pyplot as plt
import random
import math
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch import optim
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

# fix seeds
SEED = 225
random.seed(SEED)
np.random.seed(SEED)

# data processing
def f(x):
  return (x - 0.2)*(x + 0.3)*(3*x - 1.5) + 10

SIZE = 50
x = np.linspace(-2, 2, SIZE, np.float32)
y = f(x)

n_idx = random.sample(range(50), 15)

y_noise = y + 3 * np.random.randn(SIZE)

idx = list(range(SIZE))
np.random.shuffle(idx)

x_train, y_train, x_test, y_test = x[idx][:40], y_noise[idx][:40], x[idx][40:], y[idx][40:]

x_train, y_train, x_test, y_test = map(
    torch.FloatTensor, (x_train[:, np.newaxis], y_train[:, np.newaxis], x_test[:, np.newaxis], y_test[:, np.newaxis])
)
n, _ = x_train.shape

plt.plot(x_train.numpy(), y_train.detach().numpy(), 'o')
plt.plot(x_test.numpy(), y_test.detach().numpy(), 'ro')
plt.plot(x, y, '-')
x_train
def train_model(model, x_train, y_train, x_test, y_test, epochs=250, batch_size=10, lr=0.01, weight_decay=0):
  # data
  train_dataset = TensorDataset(x_train, y_train)
  train_data_loader = DataLoader(train_dataset, batch_size=batch_size)

  # loss function
  loss_func = F.mse_loss

  # optimizer
  optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)

  # training loop
  y_pred = model(x_test)
  print('Test error of model before training: ', loss_func(y_pred, y_test))

  train_loss = []
  test_loss = []
  for epoch in range(epochs):
    epoch_loss = 0
    
    for xb, yb in train_data_loader:
      
      pred = model(xb)
      loss = loss_func(pred, yb)

      loss.backward()
      optimizer.step()
      optimizer.zero_grad()
      
      epoch_loss += loss
    
    y_pred = model(x_test)
    test_loss.append(loss_func(y_pred, y_test))
    train_loss.append(epoch_loss/len(train_data_loader))
    
    print('Loss at epoch %d : %f' % (epoch, epoch_loss/len(train_data_loader)))

  # model evaluation
  y_pred = model(x_test)
  print('Test error of model after training: ', loss_func(y_pred, y_test))
  return train_loss, test_loss
def visualize_model(trained_model):
  x = torch.linspace(-2, 2, 50).unsqueeze(1)
  y_ = trained_model.forward(x)
  y = f(x)
  plt.figure()
  plt.plot(x_train.numpy(), y_train.detach().numpy(), 'o')
  plt.plot(x_test.numpy(), y_test.detach().numpy(), 'ro')
  plt.plot(x.numpy(), y_.detach().numpy(), 'b-')
  plt.plot(x.numpy(), y.detach().numpy(), 'k-')
def get_learning_curves(train_loss, test_loss):
  plt.figure()
  plt.plot(train_loss[30:])
  plt.plot(test_loss[30:])
  plt.legend(['Training loss', 'Validation loss'])<jupyter_output><empty_output><jupyter_text>The following seemingly correct model doesn't train.<jupyter_code># model definition
class mlp_regressor(nn.Module):
  def __init__(self):
      super().__init__()
      self.weights_1 = nn.Parameter(torch.randn(1, 512))
      self.bias_1 = nn.Parameter(torch.zeros(512))
      
      self.weights_2 = nn.Parameter(torch.randn(512, 512))
      self.bias_2 = nn.Parameter(torch.zeros(512))
      
      self.weights_3 = nn.Parameter(torch.randn(512, 1))
      self.bias_3 = nn.Parameter(torch.zeros(1))

  def forward(self, xb):
    h =  xb.matmul(self.weights_1) + self.bias_1
    h = F.relu(h)
    h =  h.matmul(self.weights_2) + self.bias_2
    h = F.relu(h)
    h =  h.matmul(self.weights_3) + self.bias_3
    out = h
    return out
  
model = mlp_regressor()
_, _ = train_model(model, x_train, y_train, x_test, y_test, lr=0.000001)<jupyter_output>Test error of model before training:  tensor(103720.0547, grad_fn=<MseLossBackward>)
Loss at epoch 0 : 88184.437500
Loss at epoch 1 : 2035.091553
Loss at epoch 2 : 530.837402
Loss at epoch 3 : 441.387756
Loss at epoch 4 : 401.555542
Loss at epoch 5 : 366.754852
Loss at epoch 6 : 334.945923
Loss at epoch 7 : 306.043915
Loss at epoch 8 : 279.824280
Loss at epoch 9 : 255.878998
Loss at epoch 10 : 234.181259
Loss at epoch 11 : 214.384979
Loss at epoch 12 : 196.407043
Loss at epoch 13 : 180.052704
Loss at epoch 14 : 165.137253
Loss at epoch 15 : 151.646240
Loss at epoch 16 : 139.386108
Loss at epoch 17 : 128.256699
Loss at epoch 18 : 118.166603
Loss at epoch 19 : 109.041069
Loss at epoch 20 : 100.804291
Loss at epoch 21 : 93.380280
Loss at epoch 22 : 86.673126
Loss at epoch 23 : 80.597473
Loss at epoch 24 : 75.070679
Loss at epoch 25 : 70.101349
Loss at epoch 26 : 65.618332
Loss at epoch 27 : 61.577393
Loss at epoch 28 : 57.966072
Loss at epoch 29 : 54.690826
Loss at epoch 30 : 51.729134
Lo[...]<jupyter_text>Tricks to tame a diverging model  Adjusting the learning rateWhen the loss blows up, as it is in above case, first thing to try is to decrease the learning rate. The model seems to overfit heavily to the noise. But atleast it's learning!<jupyter_code>model = mlp_regressor()
train_loss, test_loss = train_model(model, x_train, y_train, x_test, y_test, lr=0.00001)
get_learning_curves(train_loss, test_loss)
visualize_model(model)<jupyter_output>Test error of model before training:  tensor(26859.6543, grad_fn=<MseLossBackward>)
Loss at epoch 0 : 1267152.750000
Loss at epoch 1 : 222115.906250
Loss at epoch 2 : 1251.366943
Loss at epoch 3 : 890.796021
Loss at epoch 4 : 648.431824
Loss at epoch 5 : 480.726074
Loss at epoch 6 : 360.864258
Loss at epoch 7 : 273.213867
Loss at epoch 8 : 208.713684
Loss at epoch 9 : 161.165253
Loss at epoch 10 : 125.405960
Loss at epoch 11 : 98.636604
Loss at epoch 12 : 78.730026
Loss at epoch 13 : 63.903992
Loss at epoch 14 : 52.768459
Loss at epoch 15 : 44.227325
Loss at epoch 16 : 37.757172
Loss at epoch 17 : 32.775124
Loss at epoch 18 : 28.855206
Loss at epoch 19 : 25.778761
Loss at epoch 20 : 23.304943
Loss at epoch 21 : 21.254219
Loss at epoch 22 : 19.527365
Loss at epoch 23 : 18.152166
Loss at epoch 24 : 17.005974
Loss at epoch 25 : 16.018202
Loss at epoch 26 : 15.161011
Loss at epoch 27 : 14.424515
Loss at epoch 28 : 13.769180
Loss at epoch 29 : 13.183675
Loss at epoch 30 : 12.652834
Loss at [...]<jupyter_text>Initialiazation of model parametersAnother approach is to do better initialization of model parameters. Though this isn't particularly useful in this particular example, this showcases the advantage better initialization brings. Also, the overrfitting is relatively less apparent.<jupyter_code># model definition
class mlp_regressor_glorot(nn.Module):
  def __init__(self):
      super().__init__()
      self.weights_1 = nn.Parameter(torch.randn(1, 800) / math.sqrt(1))
      self.bias_1 = nn.Parameter(torch.zeros(800))
      
      self.weights_2 = nn.Parameter(torch.randn(800, 800) / math.sqrt(800))
      self.bias_2 = nn.Parameter(torch.zeros(800))
      
      self.weights_3 = nn.Parameter(torch.randn(800, 1) / math.sqrt(800))
      self.bias_3 = nn.Parameter(torch.zeros(1))

  def forward(self, xb):
    h =  xb.matmul(self.weights_1) + self.bias_1
    h = F.relu(h)
    h =  h.matmul(self.weights_2) + self.bias_2
    h = F.relu(h)
    h =  h.matmul(self.weights_3) + self.bias_3
    out = h
    return out
  
model = mlp_regressor_glorot()
train_loss, test_loss = train_model(model, x_train, y_train, x_test, y_test, lr=0.001)
get_learning_curves(train_loss, test_loss)
visualize_model(model)<jupyter_output>Test error of model before training:  tensor(77.4261, grad_fn=<MseLossBackward>)
Loss at epoch 0 : 114.131058
Loss at epoch 1 : 67.354576
Loss at epoch 2 : 67.599602
Loss at epoch 3 : 67.323547
Loss at epoch 4 : 67.020401
Loss at epoch 5 : 66.728973
Loss at epoch 6 : 66.446281
Loss at epoch 7 : 66.169205
Loss at epoch 8 : 65.915146
Loss at epoch 9 : 65.666107
Loss at epoch 10 : 65.461823
Loss at epoch 11 : 65.233551
Loss at epoch 12 : 65.014221
Loss at epoch 13 : 64.815605
Loss at epoch 14 : 64.571915
Loss at epoch 15 : 64.353119
Loss at epoch 16 : 64.139137
Loss at epoch 17 : 63.935829
Loss at epoch 18 : 63.732162
Loss at epoch 19 : 63.544262
Loss at epoch 20 : 63.363228
Loss at epoch 21 : 63.201096
Loss at epoch 22 : 63.021645
Loss at epoch 23 : 62.831833
Loss at epoch 24 : 62.667473
Loss at epoch 25 : 62.490677
Loss at epoch 26 : 62.328453
Loss at epoch 27 : 62.167740
Loss at epoch 28 : 62.028610
Loss at epoch 29 : 61.835674
Loss at epoch 30 : 61.710335
Loss at epoch 31 : 61.531380
[...]<jupyter_text>Changing the activation functionActivation function plays a vital role in the model dynamics. Changing the activation from ReLU to tanh lets the model learn in this case. But it heavily overfits!<jupyter_code># model definition
class mlp_regressor_tanh(nn.Module):
  def __init__(self):
      super().__init__()
      self.weights_1 = nn.Parameter(torch.randn(1, 512))
      self.bias_1 = nn.Parameter(torch.zeros(512))
      
      self.weights_2 = nn.Parameter(torch.randn(512, 512))
      self.bias_2 = nn.Parameter(torch.zeros(512))
      
      self.weights_3 = nn.Parameter(torch.randn(512, 1))
      self.bias_3 = nn.Parameter(torch.zeros(1))

  def forward(self, xb):
    h =  xb.matmul(self.weights_1) + self.bias_1
    h = F.tanh(h)
    h =  h.matmul(self.weights_2) + self.bias_2
    h = F.tanh(h)
    h =  h.matmul(self.weights_3) + self.bias_3
    out = h
    return out
  
model = mlp_regressor_tanh()
train_loss, test_loss = train_model(model, x_train, y_train, x_test, y_test, lr=0.001)
get_learning_curves(train_loss, test_loss)
visualize_model(model)<jupyter_output>/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1374: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.
  warnings.warn("nn.functional.tanh is deprecated. Use torch.tanh instead.")<jupyter_text>Better initialization and more suitable activation function makes our model even better<jupyter_code># model definition
class mlp_regressor_tanh_glorot(nn.Module):
  def __init__(self):
      super().__init__()
      self.weights_1 = nn.Parameter(torch.randn(1, 512) / math.sqrt(1))
      self.bias_1 = nn.Parameter(torch.zeros(512))
      
      self.weights_2 = nn.Parameter(torch.randn(512, 512) / math.sqrt(512))
      self.bias_2 = nn.Parameter(torch.zeros(512))
      
      self.weights_3 = nn.Parameter(torch.randn(512, 1) / math.sqrt(512))
      self.bias_3 = nn.Parameter(torch.zeros(1))

  def forward(self, xb):
    h =  xb.matmul(self.weights_1) + self.bias_1
    h = F.tanh(h)
    h =  h.matmul(self.weights_2) + self.bias_2
    h = F.tanh(h)
    h =  h.matmul(self.weights_3) + self.bias_3
    out = h
    return out
  
model = mlp_regressor_tanh_glorot()
train_loss, test_loss = train_model(model, x_train, y_train, x_test, y_test, lr=0.001)
get_learning_curves(train_loss, test_loss)
visualize_model(model)<jupyter_output>/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:1374: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.
  warnings.warn("nn.functional.tanh is deprecated. Use torch.tanh instead.")<jupyter_text>L2 regularizationThe model with better activation and initialization works fine, but it still seems to overfit to some extent. This can be attributed to the noise in the data and it may not always be possible to completely recover the underlying function. But anothere trick under our hats is to regularize the model weights from getting too large. This limits the complexity of the model.<jupyter_code>model = mlp_regressor_tanh_glorot()
train_loss, test_loss = train_model(model, x_train, y_train, x_test, y_test, lr=0.001, weight_decay=0.135)
get_learning_curves(train_loss, test_loss)
visualize_model(model)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11235005913161007, 'fraction_numerical': 0.10339584389254941, 'mean_word_length': 4.005919661733615, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '627009', 'license': 'MIT', 'n_tokens_mistral': 5329, 'n_tokens_neox': 4351, 'n_words': 1191, 'path': 'Lecture2_mlp_tricks.ipynb', 'repo_name': 'CharlesPoletowin/YCBS-273', 'size': '11838'}