convert.py

import numpy as np
import os
import pickle
import torch
import tqdm

def extract_metric(checkpoints: list[dict], metric_class: type, num_epochs: int) -> list:
    # Create the array, full of nans
    array = np.nan * np.ones((len(checkpoints), num_epochs))

    for i, checkpoint in enumerate(checkpoints):
        # Find the metric by type
        metric = next((v for v in checkpoint["epoch_metrics"] if isinstance(v, metric_class)), None)
        if metric is None:
            metric  = next((v for v in checkpoint["batch_metrics"] if isinstance(v, metric_class)), None)
        if metric is None:
            continue

        # Get the values & store in the array
        array[i, :len(metric.values)] = metric.values

    return array

def convert_to_npz(
        name_base: str,
        min_depth=1,
        max_depth=10,
        min_width_pow=3,
        max_width_pow=5,
        count=8,
    ):

    # Set up the models, ordered by number of parameters
    def model(depth: int, width: int):
        return (13, *([width] * depth), 1)
    specs = [model(d, 2 ** (2*w)) for d in range(min_depth, max_depth+1) for w in range(min_width_pow, max_width_pow+1)]
    specs = [*sorted(specs, key=lambda x: sum(x))]

    # Read in each of the files
    print("Reading files...")
    checkpoints = []
    metrics = set()
    n_epochs = 0
    for steps in tqdm.tqdm(specs):
        for i in range(count):
            try:
                name = name_base + " " + "-".join(str(s) for s in steps + (i,))
                checkpoint = torch.load(f"checkpoints/{name}", map_location="cpu")
                n_epochs = max(n_epochs, checkpoint["epoch"])
                checkpoints.append(checkpoint)
                for m in checkpoint["batch_metrics"]:
                    metrics.add(m.__class__)
                for m in checkpoint["epoch_metrics"]:
                    metrics.add(m.__class__)
            except FileNotFoundError:
                print(f"File not found: {name}")
    
    # Construct numpy arrays for each metric
    print("Constructing arrays...")
    metric_arrays = {}
    for metric in tqdm.tqdm(metrics):
        metric_arrays[metric.__name__] = extract_metric(checkpoints, metric, n_epochs)

    # Save the results as an NPZ file
    if not os.path.exists("results"):
        os.makedirs("results")
    print("Saving file...")
    pickle.dump(specs, open(f"results/{name_base} specs.pkl", "wb"))
    np.savez(f"results/{name_base}.npz", **metric_arrays)

if __name__ == "__main__":
    convert_to_npz("Giant House Group Adam 0")
    convert_to_npz("Giant House Group Adam 32")
    convert_to_npz("Giant House Group Adam 128")
    convert_to_npz("Giant House Group MNorm Adam 1e3 0")
    convert_to_npz("Giant House Group SGD 0")
    convert_to_npz("Giant House Group SGD 32")
    convert_to_npz("Giant House Group SGD 128")