|
|
from matplotlib import pyplot as plt |
|
|
import json |
|
|
|
|
|
if __name__ == "__main__": |
|
|
import argparse |
|
|
parser = argparse.ArgumentParser("Plot loss for the training log") |
|
|
parser.add_argument("-t", "--trainlog", required=True, help="Training log file") |
|
|
parser.add_argument("-l", "--loss", help="Loss plot to save (optional)") |
|
|
parser.add_argument("--log-scale", default=False, action="store_true", help="Log scale") |
|
|
parser.add_argument("-P", "--no-plot", default=True, action="store_false", help="Don't open matplotlib figure") |
|
|
args = parser.parse_args() |
|
|
|
|
|
train_inner_upd, train_inner_loss = [], [] |
|
|
train_upd, train_loss = [], [] |
|
|
val_upd, val_loss = [], [] |
|
|
|
|
|
with open(args.trainlog, "r") as tl: |
|
|
for line in tl: |
|
|
|
|
|
if line[0] != "{": |
|
|
continue |
|
|
try: |
|
|
data = json.loads(line.strip()) |
|
|
except: |
|
|
continue |
|
|
if "loss" in data: |
|
|
loss = float(data["loss"]) |
|
|
upd = int(data["num_updates"]) |
|
|
if len(train_inner_upd) == 0 or train_inner_upd[-1] < upd: |
|
|
train_inner_upd.append(upd) |
|
|
train_inner_loss.append(loss) |
|
|
if "valid_loss" in data: |
|
|
loss = float(data["valid_loss"]) |
|
|
upd = int(data["valid_num_updates"]) |
|
|
if len(val_upd) == 0 or val_upd[-1] < upd: |
|
|
val_upd.append(upd) |
|
|
val_loss.append(loss) |
|
|
if "train_loss" in data: |
|
|
loss = float(data["train_loss"]) |
|
|
upd = int(data["train_num_updates"]) |
|
|
if len(train_upd) == 0 or train_upd[-1] < upd: |
|
|
train_upd.append(upd) |
|
|
train_loss.append(loss) |
|
|
|
|
|
plt.figure() |
|
|
plt.plot(train_upd, train_loss, "r") |
|
|
plt.plot(val_upd, val_loss, "b") |
|
|
if len(train_inner_upd) > 0: |
|
|
plt.plot(train_inner_upd, train_inner_loss, "r", alpha=0.3) |
|
|
plt.legend(["train", "valid"]) |
|
|
if args.log_scale: |
|
|
plt.yscale("log") |
|
|
elif min(min(train_loss), min(val_loss)) < 1: |
|
|
plt.ylim((0, 1)) |
|
|
plt.xlabel("Updates") |
|
|
plt.ylabel("Loss") |
|
|
if args.loss: |
|
|
plt.savefig(args.loss) |
|
|
if args.no_plot: |
|
|
plt.show() |
|
|
|
