<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Learning the identity with neural networks
We try to reproduce the problem of learning neural networks that learn the identity function and extrapolate beyond the training range. See 1.1 (and Figure 1) in Trask et al. https
Step1:  Training and extrapolation work perfectly if we do not add any non-linearity. But adding any activation function makes extrapolation fail at some points. The first model uses RELU which is very close to linear.
Initialize model and optimizer
Step2:  Use MSE as loss function
Step3:  Generate random training data (in the range [0,1])
Step4:  Train network for 10 epochs
Step5:  Results of trained network
Step6:  Initialize second model and optimizer (closer to the paper)
<ASSISTANT_TASK:>
Python Code:
import torch
from torch import Tensor, nn, optim
import torch.nn.functional as F
import numpy as np
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt
Explanation: Learning the identity with neural networks
We try to reproduce the problem of learning neural networks that learn the identity function and extrapolate beyond the training range. See 1.1 (and Figure 1) in Trask et al. https://arxiv.org/pdf/1808.00508.pdf
End of explanation
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.h1 = nn.Linear(1, 20)
        self.h2 = nn.Linear(20, 1)
    def forward(self, xb):
        return self.h2(F.relu(self.h1(xb)))
mod = MyModel()
opt = optim.SGD(mod.parameters(), lr=.05)
mod(Tensor([7]))
Explanation: Training and extrapolation work perfectly if we do not add any non-linearity. But adding any activation function makes extrapolation fail at some points. The first model uses RELU which is very close to linear.
Initialize model and optimizer
End of explanation
loss_func = F.mse_loss
loss_func(Tensor([7]), Tensor([6]))
loss_func(mod(Tensor([7])), Tensor([7]))
Explanation: Use MSE as loss function
End of explanation
x = torch.rand(1000)
Explanation: Generate random training data (in the range [0,1])
End of explanation
for epoch in tqdm(range(10)):
    for i in x:
        pred = mod(i.unsqueeze(0))
        loss = loss_func(pred, i.unsqueeze(0))
        loss.backward()
        opt.step()
        opt.zero_grad()
Explanation: Train network for 10 epochs
End of explanation
mod(Tensor([5]))
points = [[x,mod(Tensor([x])).item()] for x in np.linspace(-10,100,1000)]
plt.scatter([x[0] for x in points], [x[1] for x in points])
Explanation: Results of trained network
End of explanation
class MyModel2(nn.Module):
    def __init__(self):
        super().__init__()
        self.h1 = nn.Linear(1, 8)
        self.h2 = nn.Linear(8, 8)
        self.h3 = nn.Linear(8, 8)
        self.h4 = nn.Linear(8, 1)
    def forward(self, xb):
        res = torch.tanh(self.h1(xb))
        res = torch.tanh(self.h2(res))
        res = torch.tanh(self.h3(res))
        res = self.h4(res)
        return res
mod2 = MyModel2()
opt2 = optim.SGD(mod2.parameters(), lr=.01)
for epoch in tqdm(range(10)):
    for i in x:
        pred = mod2(i.unsqueeze(0))
        loss = loss_func(pred, i.unsqueeze(0))
        loss.backward()
        opt2.step()
        opt2.zero_grad()
mod2(Tensor([.3]))
mod2(Tensor([-1]))
points = [[x,mod2(Tensor([x])).item()] for x in np.linspace(-2,2,1000)]
plt.scatter([x[0] for x in points], [x[1] for x in points])
plt.scatter(np.linspace(-1,1,100),np.linspace(-1,1,100))
Explanation: Initialize second model and optimizer (closer to the paper)
End of explanation
<END_TASK>