Experiments for UCSD Cogs 181 course

scramble_test.py

import numpy as np
import torch 
import torch.nn as nn
import matplotlib.pyplot as plt

def load_data(num_images=2000): 
    ''' 
    Reads in images and labels from MNIST files.
    '''
    # Reads in images
    with open('images.idx3-ubyte', 'rb') as f: 
        f.read(16)
        buff = f.read(28 * 28 * num_images)
        images = np.frombuffer(buff, dtype=np.uint8).astype(np.float32)
        images = images.reshape(num_images, 28 * 28)
    with open('labels.idx1-ubyte', 'rb') as f: 
        f.read(8)
        buff = f.read(num_images)
        labels = np.frombuffer(buff, dtype=np.uint8).astype(np.int64)
    return images, labels

def test(images, labels, net):
    with torch.no_grad():
        prediction = net(images).argmax(dim=1)
        return (prediction == labels).sum().float().item() / labels.shape[0]

def batch_iterator(images, labels, batch_size): 
    '''
    Generator to make minibatch iteration easier.
    '''
    size = labels.shape[0]
    for lower in range(0, size, batch_size): 
        upper = lower + batch_size
        yield images[lower:upper], labels[lower:upper]

def initialize_weights(layer): 
    if isinstance(layer, nn.Linear) or isinstance(layer, nn.Conv2d):
        nn.init.xavier_normal_(layer.weight)
        layer.bias.data.fill_(0.0)
    
def train_net(images, labels, net, seed, lr=0.01): 
    # Constants 
    num_ep = 100
    batch_size = 200

    np.random.seed(seed)
    torch.manual_seed(seed)
    
    # Initialize weights for reproducibility
    net.apply(initialize_weights)

    # Set up network training 
    out_activation = nn.Softmax(dim=1)
    loss = nn.NLLLoss()
    optim = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9) 

    # Set up recording 
    losses = [0] * num_ep
    acc = [0] * num_ep

    # Epochs 
    for ep in range(num_ep):
        # For mini-batches
        for img_batch, lbl_batch in batch_iterator(images, labels, batch_size):
            optim.zero_grad()
            predictions = net(img_batch)
            error = loss(predictions, lbl_batch)

            # Recording and backward
            losses[ep] += error.item()
            error.backward()
            optim.step()
        # Recording and print
        acc[ep] = test(images, labels, net)
    
    plt.plot(losses)
    plt.title('Loss over Training')
    plt.show()

    plt.plot(acc)
    plt.title('Accuracy over Training')
    plt.show()

    print(f'Lowest Loss {round(min(losses), 5)}')
    print(f'Highest Acc {round(max(acc), 5)}\n')

class shape_func(nn.Module):
    '''
    Helper class for conv nets.
    '''
    def forward(self, tensor):
        pass

def main(): 
    # Constants
    num_images = 2000
    scramble = False
    seed = 1

    # Z-score images (across single channel) 
    images, labels = load_data(num_images)
    mean, std = images.mean(), images.std()
    images = (images - mean) / std

    # Scramble images
    if scramble:
        plt.imshow((images[0] * std + mean).reshape(28, 28).astype(np.uint8))
        plt.show()
        np.random.seed(seed)
        state = np.random.get_state()
        for row in range(images.shape[0]):
            np.random.shuffle(images[row])
            np.random.set_state(state)
        plt.imshow((images[0] * std + mean).reshape(28, 28).astype(np.uint8))
        plt.show()

    # Convert to torch
    images = torch.from_numpy(images)
    labels = torch.from_numpy(labels)

    # Train linear
    print('Linear')
    net = nn.Sequential(nn.Linear(28 * 28, 10), nn.LogSoftmax(dim=1)) 
    train_net(images, labels, net, seed)

    # Reshape images for conv training
    images = images.reshape(-1, 1, 28, 28)

    # Train conv 28 sized filter
    print('Conv 28 Sized Filter')
    shape_func.forward = lambda self, tensor : tensor.reshape(-1, 10)
    net = nn.Sequential(nn.Conv2d(1, 10, 28),  shape_func(), 
        nn.LogSoftmax(dim=1))
    train_net(images, labels, net, seed)

    # Train conv 8 sized filter v1
    print('Conv 8 Sized Filter, with Linear Layer')
    shape_func.forward = lambda self, tensor : tensor.reshape(-1, 21 * 21)
    net = nn.Sequential(nn.Conv2d(1, 1, 8),  shape_func(), 
        nn.Linear(21 * 21, 10), nn.LogSoftmax(dim=1))
    train_net(images, labels, net, seed)
    # With ReLU
    print('Conv 8 Sized Filter, with Linear Layer and ReLU')
    net = nn.Sequential(nn.Conv2d(1, 1, 8), nn.ReLU(), shape_func(), 
        nn.Linear(21 * 21, 10), nn.LogSoftmax(dim=1))
    train_net(images, labels, net, seed)

    # Train conv 8 sized filter v2
    print('Fully Conv 8 Sized Filter')
    shape_func.forward = lambda self, tensor : tensor.reshape(-1, 10)
    net = nn.Sequential(nn.Conv2d(1, 1, 8),  nn.Conv2d(1, 1, 8), 
        nn.Conv2d(1, 1, 8), nn.Conv2d(1, 10, 7), shape_func(), 
        nn.LogSoftmax(dim=1))
    train_net(images, labels, net, seed, lr=0.003)
    # With ReLU
    print('Fully Conv 8 Sized Filter, with ReLU')
    net = nn.Sequential(nn.Conv2d(1, 1, 8), nn.ReLU(), nn.Conv2d(1, 1, 8), 
        nn.ReLU(), nn.Conv2d(1, 1, 8), nn.ReLU(), nn.Conv2d(1, 10, 7), 
        shape_func(), nn.LogSoftmax(dim=1))
    train_net(images, labels, net, seed, lr=0.003)

if __name__ == '__main__': 
    main()