Neural network trained to predict XOR truth table with numpy and autograd

xor_neural_network.py

import autograd.numpy as np  # Important: Use autograd's wrapped numpy
from autograd import grad
import matplotlib.pyplot as plt

# ==========================================
# 1. Dataset Preparation
# ==========================================
# Represent XOR Inputs as a Numpy array
X = np.array([
    [0,0],
    [0,1],
    [1,0],
    [1,1]
])

# represent XOR Expected Outputs as a Numpy array
y = np.array([
    [0],
    [1],
    [1],
    [0]
])


# ==========================================
# 2. Network Initialisation
# ==========================================
np.random.seed(42)

input_size = 2
hidden_size = 4
output_size = 1
learning_rate = 0.5
epochs = 10000

# Initialize weights and biases. 
# Useful: Pack them into a list so we can pass them to our loss function easily.
W1 = np.random.uniform(-1, 1, (input_size, hidden_size))
b1 = np.zeros((1, hidden_size))
W2 = np.random.uniform(-1, 1, (hidden_size, output_size))
b2 = np.zeros((1, output_size))
params = [W1, b1, W2, b2]

# ==========================================
# 3. Activation Functions
# ==========================================
def sigmoid(x):
    return 1/(1+np.exp(-x))

print(sigmoid(10))

def sigmoid_derivative(sigmoid_out):
    # Expects the output of the sigmoid function, not the raw input
    return sigmoid_out * (1 - sigmoid_out)

print(sigmoid_derivative(0.5))

# ==========================================
# 4. Forward Pass & Loss Function
# ==========================================

def predict(parameters, input):
    w_layer1, b_layer1, w_layer2, b_layer2 = parameters

    # Layer 1 (Hidden)
    z1 = np.dot(input, w_layer1) + b_layer1  # pre-activation, hidden layer
    a1 = sigmoid(z1)  # hidden layer output

    # Layer 2 (Output)
    z2 = np.dot(a1, w_layer2) + b_layer2  # pre-activation, output
    return sigmoid(z2)  # final prediction

def compute_loss(parameters):
    y_hat = predict(parameters, X)

    # Calculate Mean Squared Error Loss
    # The 0.5 factor simplifies the gradient of the loss during backpropagation.
    # np.mean() averages the loss across all 4 training samples.
    loss = np.mean(0.5 * (y_hat - y)**2)
    return loss

# Create a callable function that returns the gradients of compute_loss
# w.r.t its first argument (parameters)
calculate_gradients = grad(compute_loss)

# ==========================================
# 5. Training Loop
# ==========================================
loss_history = []

for epoch in range(epochs):
    # Record the current loss (optional, just for plotting)
    current_loss = compute_loss(params)
    loss_history.append(current_loss)

    # 1. Compute all gradients automatically
    grads = calculate_gradients(params)

    # 2. Update weights and biases using Gradient Descent
    learning_rate = 0.5

    #params[0] -= learning_rate * grads[0] # W1
    #params[1] -= learning_rate * grads[1] # b1
    #params[2] -= learning_rate * grads[2] # W2
    #params[3] -= learning_rate * grads[3] # b2

    params = [
        parameter - learning_rate * gradient
        for parameter, gradient in zip(params, grads)
    ]

    # Print progress
    if epoch % 2000 == 0:
        print(epoch, current_loss)


# ==========================================
# 6. Testing and Plotting
# ==========================================
# Print Final Predictions after Training 
for i in range(len(X)):
    prediction = predict(params, X[i])
    print(f"Input: {X[i]} | Target: {y[i][0]} | Prediction: {prediction[0][0]:.4f}")
    

# Plot the learning curve
plt.figure(figsize=(8, 5))
plt.plot(loss_history, color='green', linewidth=2)
plt.title("MLP Learning Curve for XOR (Using Autograd)")
plt.xlabel("Epochs")
plt.ylabel("Mean Squared Error (Loss)")
plt.grid(True)
plt.show()