train.py

import numpy as np
from network import NeuralNetwork

def load_mnist():
    from urllib import request
    import gzip
    
    # Import the MNIST dataset
    base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
    files = {
        "train_images": "train-images-idx3-ubyte.gz",
        "train_labels": "train-labels-idx1-ubyte.gz",
        "test_images": "t10k-images-idx3-ubyte.gz",
        "test_labels": "t10k-labels-idx1-ubyte.gz"
    }
    
    def download_and_parse_images(filename):
        print(f"Downloading {filename}...")
        response = request.urlopen(base_url + filename)
        with gzip.open(response, 'rb') as f:
            data = np.frombuffer(f.read(), np.uint8, offset=16)
        return data.reshape(-1, 784) / 255.0
    
    def download_and_parse_labels(filename):
        print(f"Downloading {filename}...")
        response = request.urlopen(base_url + filename)
        with gzip.open(response, 'rb') as f:
            return np.frombuffer(f.read(), np.uint8, offset=8)
    
    X_train = download_and_parse_images(files["train_images"])
    y_train = download_and_parse_labels(files["train_labels"])
    X_test = download_and_parse_images(files["test_images"])
    y_test = download_and_parse_labels(files["test_labels"])
    
    return X_train, y_train, X_test, y_test

if __name__ == "__main__":
    X_train, y_train, X_test, y_test = load_mnist()
    print(f"Training data: {X_train.shape}, Test data: {X_test.shape}")

    # Create validation split (10% of training data)
    val_size = int(0.1 * len(X_train))
    X_val = X_train[:val_size]
    y_val = y_train[:val_size]
    X_train_split = X_train[val_size:]
    y_train_split = y_train[val_size:]

    print(f"Training split: {X_train_split.shape}, Validation: {X_val.shape}")

    # Initialize improved network with deeper architecture
    nn = NeuralNetwork(hidden_sizes=[512, 384, 256, 128])
    nn.train(X_train_split, y_train_split, X_val, y_val,
             epochs=80, initial_lr=0.15, batch_size=128)

    # Evaluate on test set
    nn.training = False
    predictions = np.argmax(nn.forward(X_test), axis=1)
    accuracy = np.mean(predictions == y_test)
    print(f"\nFinal test accuracy: {accuracy:.4f}")