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))