plot_loss.py

import json
import argparse
import matplotlib.pyplot as plt


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("json_path", help="Path to JSON file containing loss list")
    parser.add_argument("--output", default="training_loss.png", help="Output image filename")
    args = parser.parse_args()

    with open(args.json_path, "r") as f:
        losses = json.load(f)

    iterations = list(range(1, len(losses) + 1))

    plt.plot(iterations, losses)
    plt.xlabel("Iterations")
    plt.ylabel("Training Loss")
    plt.title("Training Loss vs Iterations Dim 200")
    plt.grid(True)
    plt.tight_layout()

    plt.savefig(args.output, dpi=300)
    plt.show()


if __name__ == "__main__":
    main()