Spaces:

ProCreations
/

repro-expressivity-efficiency-hybrid-sequence

Running

repro-expressivity-efficiency-hybrid-sequence / source /official-code /micro_hf /plt_utils.py

Reproduction logbook (paper-82EJxJzG6r)

4ca4e4c verified 3 months ago

7.68 kB

import numpy as np
import torch
import matplotlib.pyplot as plt
from collections import defaultdict
import os

# colors = {"TF": {"TF": "blue", "SSM": "orange"}, "SSM": {"TF": "green", "SSM": "red"}, "TF-nC": {"TF-nC": "brown"}}
colors = {"TF-TF": "blue", "TF-SSM": "orange", "SSM-TF": "green", "SSM-SSM": "red", "TF~nC-TF~nC": "brown", "TF-TF-TF": "blue", "SSM-SSM-SSM": "red", "SSM-SSM-TF": "green"}

colors.update({"hybrid": "green", "TF": "blue", "SSM": "red"})


def Int(s): return int("".join([c for c in s if c.isnumeric()]))
def Empty(): return []



def get_val_and_bounds(data):
    mean = np.mean(data, axis=0)
    median = np.median(data, axis=0)
    
    # return mean, np.min(data, axis=0), np.max(data, axis=0)
    # return mean, mean-np.std(data, axis=0), mean+np.std(data, axis=0)
    # return median, np.min(data, axis=0), np.max(data, axis=0)
    return mean, np.quantile(data, 0.10, axis=0), np.quantile(data, 0.90, axis=0)
    # return median, np.quantile(data, 0.10, axis=0), np.quantile(data, 0.90, axis=0)


def savefig(taskname, filename):
    if "fig" not in os.listdir("results/" + taskname):
        os.mkdir("results/" + taskname + "/fig")
        
    plt.savefig("results/" + taskname + "/fig/" + filename + ".png")



split_array = ['_', 'task_name', 'layers', 'window', 'dim', 'num_heads', 'state_dim']
def plot(data, params, ind_var, diff_lines="layers", param_counts=None, x_axis=None, num_layers=2):
    fig, ax = plt.subplots()
    if x_axis == 'epochs':
        xs = defaultdict(Empty)
        ys = defaultdict(Empty)
        ys_lower = defaultdict(Empty)
        ys_upper = defaultdict(Empty)
        
        # Get the relevant data for these params
        for k in data.keys():
            d = dict(zip(split_array, k.split('_')))
            # if diff_lines != 'layer1'    and params['layer1']    != d['layer1']:         continue
            # if diff_lines != 'layer2'    and params['layer2']    != d['layer2']:         continue
            if diff_lines != 'window'    and params['window']    != Int(d['window']):    continue
            if diff_lines != 'dim'       and params['dim']       != Int(d['dim']):       continue
            if diff_lines != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
            if diff_lines != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue

            key = Int(d[diff_lines])

            xs[key] = np.arange(0, data[k].shape[1])
            ys[key], ys_lower[key], ys_upper[key] = get_val_and_bounds(data[k])

        legend = []
        keys = sorted(list(ys.keys()))
        for key in keys:
            plt.plot(xs[key], ys[key])
            legend.append(key)
        
        plt.legend(legend)

        for key in ys.keys():
            plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color='lightblue', alpha=0.08)
                

    if diff_lines == 'layers':
        xs = defaultdict(Empty)
        ys = defaultdict(Empty)
        ys_lower = defaultdict(Empty)
        ys_upper = defaultdict(Empty)

        # Get the relevant data for these params
        for k in data.keys():
            d = dict(zip(split_array, k.split('_')))
            if ind_var != 'window'    and params['window']    != Int(d['window']):    continue
            if ind_var != 'dim'       and params['dim']       != Int(d['dim']):       continue
            if ind_var != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
            if ind_var != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue

            key = d['layers']

            # print(k)
            # print(d['layers'])
            if d['layers'].split("-")[1].isnumeric() or len(d['layers'].split("-")) != num_layers:
                # print("Ignoring", key)
                continue

            if x_axis == 'params':
                xs[key].append(param_counts[k])
            else:
                xs[key].append(Int(d[ind_var]))
                
            r1, r2, r3 = get_val_and_bounds(data[k])
            ys[key].append(r1)
            ys_lower[key].append(r2)
            ys_upper[key].append(r3)

        # Sort the data so it is in order on the x axis
        for key in ys.keys():
            ys[key] = [a[1] for a in sorted(zip(xs[key], ys[key]))]
            ys_lower[key] = [a[1] for a in sorted(zip(xs[key], ys_lower[key]))]
            ys_upper[key] = [a[1] for a in sorted(zip(xs[key], ys_upper[key]))]
            xs[key].sort()

        # Plot the lines
        legend = []
        for key in ys.keys():
            if key == "SSM-SSM" and (ind_var in ["num_heads", "window"]) or key == "TF-TF" and ind_var == "state_dim":
                plt.axhline(y=np.mean(ys[key]), color=colors[key], linestyle='dashed')
                # plt.axhline(y=ys[key][0], color=colors[key], linestyle='dashed')
            else:
                plt.plot(xs[key], ys[key], c=colors[key])
            legend.append(key)
        plt.legend(legend)

        # Plot the error bars
        for key in ys.keys():
            if key == "SSM-SSM" and (ind_var in ["num_heads", "window"]) or key == "TF-TF" and ind_var == "state_dim":
                plt.fill_between(ax.get_xlim(), np.mean(ys_lower[key]), np.mean(ys_upper[key]), color=colors[key], alpha=0.08)
                # plt.fill_between(ax.get_xlim(), ys_lower[key][0], ys_upper[key][0], color=colors[key], alpha=0.08)
            else:
                plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color=colors[key], alpha=0.08)


    elif diff_lines == 'depths':
        # assert False # TODO: Doesn't plot across depth
        xs = defaultdict(Empty)
        ys = defaultdict(Empty)
        ys_lower = defaultdict(Empty)
        ys_upper = defaultdict(Empty)

        # Get the relevant data for these params
        for k in data.keys():
            d = dict(zip(split_array, k.split('_')))
            if ind_var != 'window'    and params['window']    != Int(d['window']):    continue
            if ind_var != 'dim'       and params['dim']       != Int(d['dim']):       continue
            if ind_var != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
            if ind_var != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue
            # print(d['layers'])
            if not d['layers'].split("-")[1].isnumeric(): continue
            # print("Here")

            key = d['layers'].split("-")[0] #+ "-" + d['layers'].split("-")[-1]

            if x_axis == 'params':
                xs[key].append(param_counts[key])
            else:
                xs[key].append(Int(d[ind_var]))
                
            r1, r2, r3 = get_val_and_bounds(data[k])
            ys[key].append(r1)
            ys_lower[key].append(r2)
            ys_upper[key].append(r3)

        # Sort the data so it is in order on the x axis
        for key in ys.keys():
            ys[key] = [a[1] for a in sorted(zip(xs[key], ys[key]))]
            ys_lower[key] = [a[1] for a in sorted(zip(xs[key], ys_lower[key]))]
            ys_upper[key] = [a[1] for a in sorted(zip(xs[key], ys_upper[key]))]
            xs[key].sort()

        # Plot the lines
        legend = []
        for key in ys.keys():
            plt.plot(xs[key], ys[key], c=colors[key.split("-")[0]])
            legend.append(key.split("-")[0])
        plt.legend(legend)

        # Plot the error bars
        for key in ys.keys():
            plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color=colors[key.split("-")[0]], alpha=0.08)

    if num_layers == 2:
        return diff_lines + "_" + ind_var
    if num_layers == 3:
        return diff_lines + "_" + ind_var + "_3"

	import numpy as np
	import torch
	import matplotlib.pyplot as plt
	from collections import defaultdict
	import os

	# colors = {"TF": {"TF": "blue", "SSM": "orange"}, "SSM": {"TF": "green", "SSM": "red"}, "TF-nC": {"TF-nC": "brown"}}
	colors = {"TF-TF": "blue", "TF-SSM": "orange", "SSM-TF": "green", "SSM-SSM": "red", "TF~nC-TF~nC": "brown", "TF-TF-TF": "blue", "SSM-SSM-SSM": "red", "SSM-SSM-TF": "green"}

	colors.update({"hybrid": "green", "TF": "blue", "SSM": "red"})


	def Int(s): return int("".join([c for c in s if c.isnumeric()]))
	def Empty(): return []



	def get_val_and_bounds(data):
	mean = np.mean(data, axis=0)
	median = np.median(data, axis=0)

	# return mean, np.min(data, axis=0), np.max(data, axis=0)
	# return mean, mean-np.std(data, axis=0), mean+np.std(data, axis=0)
	# return median, np.min(data, axis=0), np.max(data, axis=0)
	return mean, np.quantile(data, 0.10, axis=0), np.quantile(data, 0.90, axis=0)
	# return median, np.quantile(data, 0.10, axis=0), np.quantile(data, 0.90, axis=0)


	def savefig(taskname, filename):
	if "fig" not in os.listdir("results/" + taskname):
	os.mkdir("results/" + taskname + "/fig")

	plt.savefig("results/" + taskname + "/fig/" + filename + ".png")



	split_array = ['_', 'task_name', 'layers', 'window', 'dim', 'num_heads', 'state_dim']
	def plot(data, params, ind_var, diff_lines="layers", param_counts=None, x_axis=None, num_layers=2):
	fig, ax = plt.subplots()
	if x_axis == 'epochs':
	xs = defaultdict(Empty)
	ys = defaultdict(Empty)
	ys_lower = defaultdict(Empty)
	ys_upper = defaultdict(Empty)

	# Get the relevant data for these params
	for k in data.keys():
	d = dict(zip(split_array, k.split('_')))
	# if diff_lines != 'layer1' and params['layer1'] != d['layer1']: continue
	# if diff_lines != 'layer2' and params['layer2'] != d['layer2']: continue
	if diff_lines != 'window' and params['window'] != Int(d['window']): continue
	if diff_lines != 'dim' and params['dim'] != Int(d['dim']): continue
	if diff_lines != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
	if diff_lines != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue

	key = Int(d[diff_lines])

	xs[key] = np.arange(0, data[k].shape[1])
	ys[key], ys_lower[key], ys_upper[key] = get_val_and_bounds(data[k])

	legend = []
	keys = sorted(list(ys.keys()))
	for key in keys:
	plt.plot(xs[key], ys[key])
	legend.append(key)

	plt.legend(legend)

	for key in ys.keys():
	plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color='lightblue', alpha=0.08)


	if diff_lines == 'layers':
	xs = defaultdict(Empty)
	ys = defaultdict(Empty)
	ys_lower = defaultdict(Empty)
	ys_upper = defaultdict(Empty)

	# Get the relevant data for these params
	for k in data.keys():
	d = dict(zip(split_array, k.split('_')))
	if ind_var != 'window' and params['window'] != Int(d['window']): continue
	if ind_var != 'dim' and params['dim'] != Int(d['dim']): continue
	if ind_var != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
	if ind_var != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue

	key = d['layers']

	# print(k)
	# print(d['layers'])
	if d['layers'].split("-")[1].isnumeric() or len(d['layers'].split("-")) != num_layers:
	# print("Ignoring", key)
	continue

	if x_axis == 'params':
	xs[key].append(param_counts[k])
	else:
	xs[key].append(Int(d[ind_var]))

	r1, r2, r3 = get_val_and_bounds(data[k])
	ys[key].append(r1)
	ys_lower[key].append(r2)
	ys_upper[key].append(r3)

	# Sort the data so it is in order on the x axis
	for key in ys.keys():
	ys[key] = [a[1] for a in sorted(zip(xs[key], ys[key]))]
	ys_lower[key] = [a[1] for a in sorted(zip(xs[key], ys_lower[key]))]
	ys_upper[key] = [a[1] for a in sorted(zip(xs[key], ys_upper[key]))]
	xs[key].sort()

	# Plot the lines
	legend = []
	for key in ys.keys():
	if key == "SSM-SSM" and (ind_var in ["num_heads", "window"]) or key == "TF-TF" and ind_var == "state_dim":
	plt.axhline(y=np.mean(ys[key]), color=colors[key], linestyle='dashed')
	# plt.axhline(y=ys[key][0], color=colors[key], linestyle='dashed')
	else:
	plt.plot(xs[key], ys[key], c=colors[key])
	legend.append(key)
	plt.legend(legend)

	# Plot the error bars
	for key in ys.keys():
	if key == "SSM-SSM" and (ind_var in ["num_heads", "window"]) or key == "TF-TF" and ind_var == "state_dim":
	plt.fill_between(ax.get_xlim(), np.mean(ys_lower[key]), np.mean(ys_upper[key]), color=colors[key], alpha=0.08)
	# plt.fill_between(ax.get_xlim(), ys_lower[key][0], ys_upper[key][0], color=colors[key], alpha=0.08)
	else:
	plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color=colors[key], alpha=0.08)


	elif diff_lines == 'depths':
	# assert False # TODO: Doesn't plot across depth
	xs = defaultdict(Empty)
	ys = defaultdict(Empty)
	ys_lower = defaultdict(Empty)
	ys_upper = defaultdict(Empty)

	# Get the relevant data for these params
	for k in data.keys():
	d = dict(zip(split_array, k.split('_')))
	if ind_var != 'window' and params['window'] != Int(d['window']): continue
	if ind_var != 'dim' and params['dim'] != Int(d['dim']): continue
	if ind_var != 'num_heads' and params['num_heads'] != Int(d['num_heads']): continue
	if ind_var != 'state_dim' and params['state_dim'] != Int(d['state_dim']): continue
	# print(d['layers'])
	if not d['layers'].split("-")[1].isnumeric(): continue
	# print("Here")

	key = d['layers'].split("-")[0] #+ "-" + d['layers'].split("-")[-1]

	if x_axis == 'params':
	xs[key].append(param_counts[key])
	else:
	xs[key].append(Int(d[ind_var]))

	r1, r2, r3 = get_val_and_bounds(data[k])
	ys[key].append(r1)
	ys_lower[key].append(r2)
	ys_upper[key].append(r3)

	# Sort the data so it is in order on the x axis
	for key in ys.keys():
	ys[key] = [a[1] for a in sorted(zip(xs[key], ys[key]))]
	ys_lower[key] = [a[1] for a in sorted(zip(xs[key], ys_lower[key]))]
	ys_upper[key] = [a[1] for a in sorted(zip(xs[key], ys_upper[key]))]
	xs[key].sort()

	# Plot the lines
	legend = []
	for key in ys.keys():
	plt.plot(xs[key], ys[key], c=colors[key.split("-")[0]])
	legend.append(key.split("-")[0])
	plt.legend(legend)

	# Plot the error bars
	for key in ys.keys():
	plt.fill_between(xs[key], ys_lower[key], ys_upper[key], color=colors[key.split("-")[0]], alpha=0.08)

	if num_layers == 2:
	return diff_lines + "_" + ind_var
	if num_layers == 3:
	return diff_lines + "_" + ind_var + "_3"