a NN visualization experiment

relu_viz.py

from PIL import Image, ImageOps, ImageDraw
import numpy as np
import pandas as pd
import os, sys
#import ace_tools as tools

# Function to load an image from a file
def load_image(file_path):
    return Image.open(file_path)

# Step 2: Grayscale the PNG file
def grayscale_image(img):
    gray_img = ImageOps.grayscale(img)
    return gray_img

# Step 3: Truncate the PNG file so it is square
def truncate_image(img):
    min_side = min(img.size)
    left = (img.width - min_side) // 2
    top = (img.height - min_side) // 2
    right = (img.width + min_side) // 2
    bottom = (img.height + min_side) // 2
    square_img = img.crop((left, top, right, bottom))
    return square_img

# Step 4: Convert the PNG file to (x, y, z) triples
def image_to_triples(img):
    img = np.array(img)
    height, width = img.shape
    triples = []
    for y in range(height):
        for x in range(width):
            z = img[y, x] / 255.0
            triples.append((x / width, y / height, z))
    return np.array(triples)

# Load and process the image

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--inputimage", help="Input PNG image to approximate", type=str,
                    default = "/home/thomasdullien/Downloads/black_circle.png")
parser.add_argument("--first_layer_neurons", help="How many neurons in the first layer",
                    type=int, default=10)
parser.add_argument("--draw_red_lines", action=argparse.BooleanOptionalAction)

args = parser.parse_args()
file_path = args.inputimage
LAYER1_NEURONS = args.first_layer_neurons

filename = os.path.split(file_path)[1]
img = load_image(file_path)
grayscale_img = grayscale_image(img)
square_img = truncate_image(grayscale_img)
triples = image_to_triples(square_img)

# Display the first few triples for verification
df_triples = pd.DataFrame(triples, columns=['x', 'y', 'z'])

import torch
import torch.nn as nn
import torch.optim as optim

# Step 5: Create and train a 1-layer ReLU network
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc = nn.Linear(2, LAYER1_NEURONS)
        self.relu = nn.ReLU()
        self.out = nn.Linear(LAYER1_NEURONS, 1)
    
    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        x = self.out(x)
        return x

def train_network(triples):
    model = SimpleNN()
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    x = torch.tensor(triples[:, :2], dtype=torch.float32)
    y = torch.tensor(triples[:, 2], dtype=torch.float32).unsqueeze(1)

    for epoch in range(10000):
        optimizer.zero_grad()
        outputs = model(x)
        loss = criterion(outputs, y)
        loss.backward()
        optimizer.step()
        print("Epoch %d: Training loss is now %f" % (epoch, loss))
        if epoch % 100 == 0:
            write_model_and_decision_boundaries(model, epoch, loss)
    return model

def write_model_and_decision_boundaries(model, train_step, loss):
    print("Writing the model and polytopes: Beginning calculation of derivatives.")
    derivative_image, value_image = calculate_derivative(model)
    print("Creating new image.")
    create_new_image(derivative_image, value_image, "./%s-%d-step-%04.04d.png" % (filename, LAYER1_NEURONS, train_step), loss)


# Step 6: Calculate the derivative of the model
def calculate_derivative(model):
    points = np.linspace(0, 1, 301)
    derivative_image = np.zeros((301, 301))
    value_image = np.zeros((301, 301))
    print("calculating 300 rows of derivatives...")
    for i, x in enumerate(points):
        print(".", end='')
        for j, y in enumerate(points):
            input_tensor = torch.tensor([[x, y]], dtype=torch.float32)
            input_tensor.requires_grad = True
            output = model(input_tensor)
            value_image[i, j] = output
            output.backward()
            gradients = input_tensor.grad.numpy()
            derivative_image[i, j] = gradients[0, 0]**2 + gradients[0, 1]**2
    return (derivative_image, value_image)

# Step 7: Create a new PNG file from the data points
def create_new_image(derivative_image, value_image, filename, loss):
    new_image = Image.new("RGB", (301, 301))
    new_image_clean = Image.new("RGB", (301, 301))
    for x in range(300):
        for y in range(300):
            pixel_color = int(value_image[x,y] * 255)
            new_image_clean.putpixel((x, y), (pixel_color, pixel_color, pixel_color))
            if derivative_image[x, y] != derivative_image[x+1, y] or derivative_image[x, y] != derivative_image[x, y+1]:
                new_image.putpixel((x, y), (255, 0, 0))
            else:
                new_image.putpixel((x, y), (pixel_color, pixel_color, pixel_color))
    new_image.save(filename)
    #new_image_clean.save(filename+"clean")

# Train the network and calculate the derivative image
print("About to train the network.")
model = train_network(triples)
#print("Done training the network. Beginning calculation of derivatives.")
#derivative_image, value_image = calculate_derivative(model)
#print("Creating new image.")
#create_new_image(derivative_image, value_image)