-
Notifications
You must be signed in to change notification settings - Fork 0
/
plot_grid.py
executable file
·63 lines (47 loc) · 1.66 KB
/
plot_grid.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
"""Plot the learning rate vs. weight decay grid."""
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle
import re
import argparse
SAVE = os.getenv("CACHED")
assert os.path.isdir(str(SAVE)), f"Could not find data folder: {SAVE}"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--optim", default="SGD", choices=["SGD", "AdamW"])
parser.add_argument("--data_dir", type=str, default="grid-norms")
return parser.parse_args()
args = parse_args()
lrs = []
wds = []
dnorms = []
data_path = f"{SAVE}/{args.data_dir}"
for name in os.listdir(data_path):
result = re.match(r"(.*)\-lr\=(.*)\-wd\=(.*)\.dat", name)
optim = result[1]
lr = float(result[2])
wd = float(result[3])
with open(os.path.join(data_path, name), "rb") as fh:
norms = pickle.load(fh)
dnorm = norms[-1] - norms[0]
if optim == args.optim:
lrs.append(lr)
wds.append(wd)
# TODO: Figure this out??
# Alternative could plot 0/1.
log_dnorm = np.sign(dnorm) * np.log(np.abs(dnorm) + 1)
dnorms.append(log_dnorm)
if not os.path.isdir("figs/grid"):
os.makedirs("figs/grid")
lrs, wds, dnorms = np.array(lrs), np.array(wds), np.array(dnorms)
plt.contourf(lrs.reshape(20, 20), wds.reshape(20, 20), dnorms.reshape(20, 20), levels=2)
# sc = plt.scatter(lrs, wds, c=dnorms)
# plt.colorbar(sc)
plt.title(fR"Log norm growth with {args.optim} after $1$ epoch by $\eta, \lambda$")
plt.xlabel(R"$\eta$")
plt.ylabel(R"$\lambda$")
plt.xscale("log")
plt.yscale("log")
plt.savefig(f"figs/grid/{args.optim}.pdf")
# contourf, pcolormesh?