import torch import torch.nn as nn import torch.nn.functional as F import torch.distributions as D from torch.distributions.multivariate_normal import MultivariateNormal import matplotlib.pyplot as plt import numpy as np

class MyNetwork(nn.Module): def __init__(self, num_input_nodes=5, num_hidden_nodes=10, num_output_nodes=1, activation_function="ReLU"): super(MyNetwork, self).__init__() # Initialize parameters and create a forward function here: self.W = nn.Parameter(0.1 * torch.randn(num_input_nodes, num_hidden_nodes)) # 0.1: Divide by num_hidden_nodes self.b = nn.Parameter(0.1 * torch.randn(num_hidden_nodes)) self.c = nn.Parameter(0.1 * torch.randn(num_hidden_nodes, num_output_nodes)) if activation_function == "ReLU": self.activation = nn.ReLU() else: self.activation = nn.Sigmoid() def forward(self, x): # [B, num_input_nodes] x = self.activation(x @ self.W + self.b) x = x @ self.c # [B, num_output_nodes] return x

# Visualizing our target function def target_function(x,y): return (np.sin(x) + np.sin(y)) x = np.linspace(-np.pi,np.pi,100) y = np.linspace(-np.pi,np.pi,100) X, Y = np.meshgrid(x, y) Z = target_function(X,Y) fig, ax = plt.subplots() ax.contourf(X, Y, Z, levels = np.linspace(-2, 2, 10)) fig.show()

# Set size of dataset and relative size of train set & test set num_examples = 5000 train_proportion = 0.8 # Sample uniformly from our area of interest x = np.random.uniform(-np.pi, np.pi, [num_examples, 2]) y = target_function(x[:,0], x[:,1]) # Convert to tensors x = torch.tensor(x) y = torch.tensor(y) # Make the train/test split num_train = int(num_examples * .8) num_test = num_examples - num_train random_indices = torch.randperm(num_examples) train_index = random_indices[0:num_train] test_index = random_indices[num_train:]

# Implement your learning algorithm here: model = MyNetwork(num_input_nodes=2, num_hidden_nodes=20, num_output_nodes=1, activation_function="Sigmoid") batch_size = 1 epoch = 50 lr = 1e-3 num_train_batch = (num_train + batch_size -1) // batch_size num_test_batch = (num_test + batch_size -1) // batch_size train_losses, test_losses = [], [] for _ in range(epoch): train_order = train_index[torch.randperm(num_train)] epoch_train_loss = 0 epoch_test_loss = 0 for batch in range(num_train_batch): indices = train_order[batch * batch_size:(batch + 1) * batch_size] x_batch = x[indices] y_batch = y[indices] pred = model(x_batch.float()) loss = ((pred - y_batch) ** 2).mean() model.zero_grad() loss.backward() for p in model.parameters(): with torch.no_grad(): p.copy_(p - lr * p.grad) epoch_train_loss += loss.item() for batch in range(num_test_batch): indices = test_index[batch * batch_size:(batch + 1) * batch_size] x_batch = x[indices] y_batch = y[indices] with torch.no_grad(): pred = model(x_batch.float()) loss = ((pred - y_batch) ** 2).mean() epoch_test_loss += loss.item() train_losses.append(epoch_train_loss) test_losses.append(epoch_test_loss) epochs = range(1, 51) fig, ax = plt.subplots() ax.plot(epochs, train_losses) ax.plot(epochs, test_losses) ax.set_xlabel('epochs') ax.set_ylabel('loss') ax.set_title('Training MyNetwork') ax.legend(['Train loss', 'Test loss']) fig.show()

# Generate a 2D contour plot of your best solution here: x = torch.linspace(-np.pi, np.pi, 100) y = torch.linspace(-np.pi, np.pi, 100) X, Y = torch.meshgrid(x, y) griddata = (torch.stack((torch.flatten(X), torch.flatten(Y))).T) Z = np.reshape(model(griddata).detach().numpy(), [100, 100]) plt.contourf(X.detach().numpy(), Y.detach().numpy(), Z, levels = np.linspace(-2, 2, 10))