part1-1.py

# -*- coding: utf-8 -*-
"""chaitanya 1_a(edited plots).ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1JTbz3KulgcULGfaidbiX3Ao9caRdjmIl
"""

pip install tensorflow==2.4

# This program trains three feedforward DNNs on two different functions: cos(x) and arcsinh(5*x*np.pi)
# Graphs are produces showing loss of all model during training and superimposes groudtruth over predictions

import tensorflow as tf
import numpy as np
import torch
import torchvision as tv
from torchvision import transforms, datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')

# Create random data between (-10, 10) and determine groundtruth
# Ground Truth Cos Function - y = cos(x)
# Ground Truth Sin Function - y = arcsinh(5*np.pi*x)
simulatedInput = 20 * torch.rand((3000, 1)) - 10
groundTruthCos = np.cos(simulatedInput)
groundTruthSin = np.arcsinh(5*np.pi*simulatedInput)

# Calculate the number of parameters in a neural network
def calcParams(inputModel):
    val = sum(params.numel() for params in inputModel.parameters() if params.requires_grad)
    return val

# Neural Network Definitions
# Shallow NN for simulation - 1 Hidden layer / 901 Parameters
class ShallowSimNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 300)
        self.fc2 = nn.Linear(300, 1)

    def forward(self, val):
        val = F.relu(self.fc1(val))
        val = self.fc2(val)
        return val

# Middle NN for simulation - 5 Hidden layer / 904 Parameters
class MiddleSimNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 25)
        self.fc2 = nn.Linear(25, 20)
        self.fc3 = nn.Linear(20, 13)
        self.fc4 = nn.Linear(13, 4)
        self.fc5 = nn.Linear(4, 1)

    def forward(self, val):
        val = F.relu(self.fc1(val))
        val = F.relu(self.fc2(val))
        val = F.relu(self.fc3(val))
        val = F.relu(self.fc4(val))
        val = self.fc5(val)
        return val    
    
# Deep NN for simulation -  5 Hidden Layers / 904 Parameters
class DeepSimNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 26)
        self.fc2 = nn.Linear(26, 19)
        self.fc3 = nn.Linear(19, 10)
        self.fc4 = nn.Linear(10, 7)        
        self.fc5 = nn.Linear(7, 5)
        self.fc6 = nn.Linear(5, 3)
        self.fc7 = nn.Linear(3, 1)        
        
    def forward(self, val):
        val = F.relu(self.fc1(val))
        val = F.relu(self.fc2(val))
        val = F.relu(self.fc3(val))
        val = F.relu(self.fc4(val))
        val = F.relu(self.fc5(val))
        val = F.relu(self.fc6(val))
        val = self.fc7(val)
        return val

model1 = ShallowSimNN()
print(calcParams(model1))
model2=MiddleSimNN()
print(calcParams(model2))
model3=DeepSimNN()
print(calcParams(model3))

# Set up necessary auxilaries for neural net training for simulation of both functions
shallowCosNN = ShallowSimNN()
middleCosNN = MiddleSimNN()
deepCosNN = DeepSimNN()
shallowSinNN = ShallowSimNN()
middleSinNN = MiddleSimNN()
deepSinNN = DeepSimNN()
costFunc = nn.MSELoss()
shallowCosOpt = optim.Adam(shallowCosNN.parameters(), lr=0.001)
middleCosOpt = optim.Adam(middleCosNN.parameters(), lr=0.001)
deepCosOpt = optim.Adam(deepCosNN.parameters(), lr=0.001)
shallowSinOpt = optim.Adam(shallowSinNN.parameters(), lr=0.001)
middleSinOpt = optim.Adam(middleSinNN.parameters(), lr=0.001)
deepSinOpt = optim.Adam(deepSinNN.parameters(), lr=0.001)

# Train neural networks and track progression on function cos(x)
EPOCHS = 5000
counter = 0
counterList = []
shallowCosCostList = []
for index in range(EPOCHS):
    counterList.append(counter)
    counter += 1
    shallowCosNN.zero_grad()
    output = shallowCosNN(simulatedInput)
    cost = costFunc(output, groundTruthCos)
    shallowCosCostList.append(cost)
    cost.backward()
    shallowCosOpt.step()

middleCosCostList = []
for index in range(EPOCHS):
    middleCosNN.zero_grad()
    output = middleCosNN(simulatedInput)
    cost = costFunc(output, groundTruthCos)
    middleCosCostList.append(cost)
    cost.backward()
    middleCosOpt.step()    
    
deepCosCostList = []
for index in range(EPOCHS):
    deepCosNN.zero_grad()
    output = deepCosNN(simulatedInput)
    cost = costFunc(output, groundTruthCos)
    deepCosCostList.append(cost)
    cost.backward()
    deepCosOpt.step()

# Train neural networks and track progression on function arcsinh(5*np.pi*x)
shallowSinCostList = []
for index in range(EPOCHS):
    shallowSinNN.zero_grad()
    output = shallowSinNN(simulatedInput)
    cost = costFunc(output, groundTruthSin)
    shallowSinCostList.append(cost)
    cost.backward()
    shallowSinOpt.step()

middleSinCostList = []
for index in range(EPOCHS):
    middleSinNN.zero_grad()
    output = middleSinNN(simulatedInput)
    cost = costFunc(output, groundTruthSin)
    middleSinCostList.append(cost)
    cost.backward()
    middleSinOpt.step()    
    
deepSinCostList = []
for index in range(EPOCHS):
    deepSinNN.zero_grad()
    output = deepSinNN(simulatedInput)
    cost = costFunc(output, groundTruthSin)
    deepSinCostList.append(cost)
    cost.backward()
    deepSinOpt.step()

# Visulaize Training process of cos function
plt.plot(counterList, shallowCosCostList, 'c', label='Shallow-1 hidden layer')
plt.plot(counterList, middleCosCostList, 'y', label='Middle-4 hidden layers')
plt.plot(counterList, deepCosCostList, 'r', label='Deep-6 hidden layers')
plt.title("Learning Progression for cos(x)")
plt.xlabel("EPOCHS")
plt.ylabel("Mean Squared Error")
plt.legend(loc="upper right")
plt.show()

# Visulaize Training process of arcsinh(x) function
plt.plot(counterList, shallowSinCostList, 'c', label='Shallow-1 hidden layer')
plt.plot(counterList, middleSinCostList, 'y', label='Middle-4 hidden layer')
plt.plot(counterList, deepSinCostList, 'r', label='Deep-6 hidden layer')
plt.title("Learning Progression for arcsinh(5*np.pi*x)")
plt.xlabel("EPOCHS")
plt.ylabel("Mean Squared Error")
plt.legend(loc="upper right")
plt.show()

# Visulaize how well the neural nets do with unseen data
simulatedInput = 20 * torch.rand((3000, 1)) - 10
groundTruthCos = np.cos(simulatedInput)
groundTruthSin = np.arcsinh(5*np.pi*simulatedInput)

# run output through nns and get predictions
shallowCosOutput = shallowCosNN(simulatedInput)
middleCosOutput = middleCosNN(simulatedInput)
deepCosOutput = deepCosNN(simulatedInput)
shallowSinOutput = shallowSinNN(simulatedInput)
middleSinOutput = middleSinNN(simulatedInput)
deepSinOutput = deepSinNN(simulatedInput)

# plot predictions for abscos(x) and compate to ground truth
plt.plot(simulatedInput, shallowCosOutput.tolist(), 'c.', label='Shallow-1 hidden layer')
plt.plot(simulatedInput, middleCosOutput.tolist(), 'y.', label='Middle-4 hidden layers')
plt.plot(simulatedInput, deepCosOutput.tolist(), 'r.', label='Deep-6 hidden layers')
plt.plot(simulatedInput, groundTruthCos.tolist(), 'b.', label='Ground Truth')
plt.title("Predictions of Models for cos(x)")
plt.xlabel("Input values")
plt.ylabel("Output values")
plt.legend(loc="lower right")
plt.show()

# plot predictions for abscos(5*np.pi*x) and compate to ground truth
plt.plot(simulatedInput, shallowSinOutput.tolist(), 'c.', label='Shallow-1')
plt.plot(simulatedInput, middleSinOutput.tolist(), 'y.', label='Middle-4')
plt.plot(simulatedInput, deepSinOutput.tolist(), 'r.', label='Deep-6')
plt.plot(simulatedInput, groundTruthSin.tolist(), 'b.', label='Ground Truth')
plt.title("Predictions of Models for arcsinh(5*np.pi*x)")
plt.xlabel("Input values")
plt.ylabel("Output values")
plt.legend(loc="upper left")
plt.show()