plot.py

import torch
import torch.nn as nn
import torch.nn.functional as F

import matplotlib.pyplot as plt
from approximations import *

import os
import argparse
import json


def json_fetcher(file_name, property="acc"):
    with open(file_name) as file:
        acc = json.load(file)[property]
        file.close()
        return acc


def sort_file_names_by_test_acc(file_names):
    l = [(file_name, json_fetcher(file_name, "acc")) for file_name in file_names]
    comp = lambda pair: pair[1]
    l.sort(key=comp)
    l.reverse()
    return l


def plot_functions(ax, fts, xbounds, numPoints=5000):
    # for each f \in fts, plots f(x) for x \in xbounds, with numPoints points
    xs, xf = xbounds
    x = torch.tensor([xs + i * (xf - xs) / numPoints for i in range(numPoints + 1)])
    for label, f, color in fts:
        y = f(x)
        ax.plot(x, y, color, label=label)


def plot_acc_for_approx(ax, poly_degrees, file_namer, label=None, color=None):
    file_names = list(map(file_namer, poly_degrees))
    test_accs = list(map(json_fetcher, file_names))
    ax.plot(poly_degrees, test_accs, color, label=label)


def plot_total_profiles(ax, file_name, label=None, color=None):
    totals = json_fetcher(file_name, "total")
    ax.plot(list(range(len(totals))), totals, color, label=label)


def plot_layerwise_profiles(ax, file_namer, label=None, color=None):
    # TODO implement this
    pass


if __name__ == "__main__":
    # fmt: off
    parser = argparse.ArgumentParser(description="training CIFAR-10")
    parser.add_argument("--save_dir", default="./plots", help="plotted images path")
    parser.add_argument("--log_dir", default="./logs", help="log path that is generated by ./all_eval.sh")
    # fmt: on

    args = parser.parse_args()

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    # chebyshev.png: Chebyshev's absolute function approximation polynomials
    plt.clf()
    _, ax = plt.subplots()
    plot_functions(
        ax,
        fts=[
            ("Abs", lambda x: abs_block(x, -1, "naive"), "b"),
            ("Chebyshev 5", lambda x: abs_block(x, 5, "chebyshev"), "g--"),
            ("Chebyshev 10", lambda x: abs_block(x, 10, "chebyshev"), "c--"),
            ("Chebyshev 15", lambda x: abs_block(x, 15, "chebyshev"), "r--"),
        ],
        xbounds=(-1.1, 1.1),
    )
    ax.set_ylim(-0.1, 1.1)
    ax.set_title("Chebyshev Approximation for Abs")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "chebyshev.png"))

    # remez.png: Remez's absolute function approximation polynomials
    plt.clf()
    _, ax = plt.subplots()
    plot_functions(
        ax,
        fts=[
            ("Abs", lambda x: abs_block(x, -1, "naive"), "b"),
            ("Remez 5", lambda x: abs_block(x, 5, "remez"), "g--"),
            ("Remez 10", lambda x: abs_block(x, 10, "remez"), "c--"),
            ("Remez 15", lambda x: abs_block(x, 15, "remez"), "r--"),
        ],
        xbounds=(-1.1, 1.1),
    )
    ax.set_ylim(-0.1, 1.1)
    ax.set_title("Remez's Algorithm for Abs")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "remez.png"))

    # relu.png: Approximation polynomials of ReLU
    plt.clf()
    _, ax = plt.subplots()
    plot_functions(
        ax,
        fts=[
            ("ReLU", lambda x: relu_block(x, -1, "naive"), "b"),
            ("Chebyshev 7", lambda x: relu_block(x, 7, "chebyshev"), "g--"),
            ("Remez 7", lambda x: relu_block(x, 7, "remez"), "r--"),
        ],
        xbounds=(-1.5, 1.5),
    )
    ax.set_ylim(-0.1, 1.1)
    ax.set_title("ReLU Approximations")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "relu.png"))

    # leaky.png: Affects of leaky method
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev L1",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez L1",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev Leaky L1",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez Leaky L1",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.5, 1.05)
    ax.set_title("Affects of Leaky ReLU Option for L1 Act Decay")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "leaky.png"))

    # leaky.png: Affects of leaky method
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l2_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev L2",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l2_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez L2",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l2_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev Leaky L2",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l2_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez Leaky L2",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.25, 1.05)
    ax.set_title("Affects of Leaky ReLU Option for L2 Act Decay")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "leaky_l2.png"))

    # norm.png: Affects of norms for activation decays
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev L1",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez L1",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l2_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev Leaky L2",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_leaky_l2_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez Leaky L2",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.5, 1.05)
    ax.set_title("Affects of Activation Decay Norms")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "norm.png"))

    # decay.png: Affects of activation decays' strengths
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev 1e-2",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez 1e-2",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-3.json"
        ),
        label="Chebyshev 1e-3",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-3.json"
        ),
        label="Remez 1e-3",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.5, 1.05)
    ax.set_title("Affects of Final Activation Decay Strengths")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "decay.png"))

    # finetune.png: Affects of finetuning stage
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev Fine-tuned",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez Fine-tuned",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_chebyshev_{x}.json"
        ),
        label="Chebyshev Just Replaced",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_remez_{x}.json"
        ),
        label="Remez Just Replaced",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.05, 1.05)
    ax.set_title("Fine-tuning Phase v.s. Just Switching into Approx")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "finetune.png"))

    # act.png: Affects of large activatin decay stage
    plt.clf()
    _, ax = plt.subplots()
    poly_degrees = list(range(5, 15 + 1, 1))
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev + Large Act Phase",
        color="bo-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_act_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez + Large Act Phase",
        color="ro-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_tuned_chebyshev_{x}_1e-2.json"
        ),
        label="Chebyshev w/o Large Act Phase",
        color="go-",
    )
    plot_acc_for_approx(
        ax,
        poly_degrees=poly_degrees,
        file_namer=lambda x: os.path.join(
            args.log_dir, f"simple_mnist_l1_tuned_remez_{x}_1e-2.json"
        ),
        label="Remez w/o Large Act Phase",
        color="mo-",
    )
    ax.set_xticks(poly_degrees)
    ax.set_ylim(0.05, 1.05)
    ax.set_title("Fine-tuning Phase v.s. Just Switching into Approx")
    ax.set_title("Affect of Large Act Decay Phase before Fine-tuning")
    ax.set_xlabel("poly degree")
    ax.set_ylabel("test acc")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "act.png"))

    # finetune_profiles.png: Activation magnitude profiles for naive vs finetuned
    plt.clf()
    _, ax = plt.subplots()
    xticklabels = [f"<2^{n}" for n in range(-4, 9)] + [">2^8"]
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_l1_act_tuned_chebyshev_10_1e-2.json"
        ),
        label="Chebyshev deg 10 Fine-tuned",
        color="bo-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_l1_act_tuned_remez_10_1e-2.json"
        ),
        label="Remez deg 10 Fine-tuned",
        color="ro-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(args.log_dir, "simple_mnist_l1_chebyshev_10.json"),
        label="Chebyshev deg 10 Just Replaced",
        color="go-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(args.log_dir, "simple_mnist_l1_remez_10.json"),
        label="Remez deg 10 Just Replaced",
        color="mo-",
    )
    ax.set_xticks(list(range(len(xticklabels))))
    ax.set_xticklabels(xticklabels, fontsize=8, rotation=10)
    ax.set_yscale("symlog")
    ax.set_ybound(-1, 10 ** 9)
    ax.set_title("Fine-tuning Phase v.s. Just Switching into Approx")
    ax.set_title("Affects of Fine-tuning for Activation Magnitude Occurences")
    ax.set_xlabel("activation magnitude")
    ax.set_ylabel("occurences")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "finetune_profile.png"))

    all_log_file_names = [
        os.path.join(args.log_dir, log) for log in os.listdir(args.log_dir)
    ]
    sorted_test_acc = sort_file_names_by_test_acc(all_log_file_names)
    print(sorted_test_acc[:20])

    # tops_profiles.png: Plotting top-ranked profiles
    plt.clf()
    _, ax = plt.subplots()
    xticklabels = [f"<2^{n}" for n in range(-4, 9)] + [">2^8"]
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_leaky_l2_act_tuned_chebyshev_7_1e-2.json"
        ),
        label="Chebyshev Leaky L2(0.01) deg 7 (acc: 0.9871)",
        color="bo-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_l1_act_tuned_chebyshev_7_1e-3.json"
        ),
        label="Chebyshev L1(0.001) deg 7 (acc: 0.9839)",
        color="ro-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_leaky_l2_act_tuned_remez_7_1e-2.json"
        ),
        label="Remez Leaky L2(0.01) deg 7 (acc: 0.9821)",
        color="go-",
    )
    plot_total_profiles(
        ax,
        file_name=os.path.join(
            args.log_dir, "simple_mnist_l1_act_tuned_remez_6_1e-3.json"
        ),
        label="Remez L1(0.001) deg 6 (acc: 0.9811)",
        color="mo-",
    )
    ax.set_xticks(list(range(len(xticklabels))))
    ax.set_xticklabels(xticklabels, fontsize=8, rotation=10)
    ax.set_yscale("symlog")
    ax.set_ybound(-1, 10 ** 9)
    ax.set_title("Fine-tuning Phase v.s. Just Switching into Approx")
    ax.set_title("Activation Magnitude Occurences of Top Accuracies Configs")
    ax.set_xlabel("activation magnitude")
    ax.set_ylabel("occurences")
    ax.legend()
    plt.savefig(os.path.join(args.save_dir, "tops_profile.png"))