-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathplot_losses.py
More file actions
73 lines (66 loc) · 2.07 KB
/
Copy pathplot_losses.py
File metadata and controls
73 lines (66 loc) · 2.07 KB
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
64
65
66
67
68
69
70
71
72
73
import argparse
import json
from pathlib import Path
import matplotlib.pyplot as plt
def main():
DESCRIPTION = """Plots loss from PyTorch training."""
parser = argparse.ArgumentParser(description=DESCRIPTION)
parser.add_argument(
"dirs", help="Path to the training run directory", type=Path, nargs="+"
)
parser.add_argument(
"-log-x", help="Plot the x-axis on a logarithmic scale", action="store_true"
)
parser.add_argument(
"-log-y", help="Plot the y-axis on a logarithmic scale", action="store_true"
)
parser.add_argument(
"-y-min",
help="Minimum value on the y-axis",
type=float,
)
parser.add_argument(
"-y-max",
help="Maximum value on the y-axis",
type=float,
)
parser.add_argument(
"-loss",
help="Which loss to plot, e.g., p0loss for policy loss. Defaults to total loss",
default="loss",
)
parser.add_argument(
"-output", help="Path to write the loss plot", type=Path, required=True
)
parser.add_argument(
"-validation",
help="Print validation loss instead of train loss",
action="store_true",
)
args = parser.parse_args()
for d in args.dirs:
file_modifier = "val" if args.validation else "train"
metrics_file = d / "train" / "t0" / f"metrics_{file_modifier}.json"
nsamp_key = "nsamp_train" if args.validation else "nsamp"
nsamps = []
losses = []
with open(metrics_file) as f:
for line in f:
metrics = json.loads(line)
nsamps.append(metrics[nsamp_key])
losses.append(metrics[args.loss])
plt.plot(nsamps, losses, label=d.name)
plt.xlabel("steps")
plt.ylabel(args.loss)
if args.y_min is not None:
plt.ylim(args.y_min, None)
if args.y_max is not None:
plt.ylim(None, args.y_max)
if args.log_x:
plt.xscale("log")
if args.log_y:
plt.yscale("log")
plt.legend()
plt.savefig(args.output)
if __name__ == "__main__":
main()