aqacf-ml-main/firealarm_ml/plot_learning_curves.py

#!/usr/bin/env python
import os
import click
from glob import glob
from tqdm import tqdm
from pathlib import Path
import numpy as np
import torch
import matplotlib.pyplot as plt


@click.command()
@click.argument('inputdir', type=click.Path(
    exists=True, dir_okay=True, file_okay=False, path_type=Path
))
def main(inputdir):

    ptfiles = glob(os.path.join(inputdir, '*.pt'))

    train_loss = {}
    test_loss = {}
    for ptfile in tqdm(ptfiles):
        checkpoint = torch.load(ptfile, map_location=torch.device('cpu'))
        epoch = checkpoint['epoch']
        train_loss[epoch] = checkpoint['train_loss']
        test_loss[epoch] = (
            checkpoint['validation_loss'] if 'validation_loss' in checkpoint
            else checkpoint['test_loss']
        )

    epochs = np.array(sorted(train_loss.keys()))
    train_loss_arr = np.array([
        train_loss[e] for e in epochs
    ])
    test_loss_arr = np.array([
        test_loss[e] for e in epochs
    ])

    best_epoch = epochs[np.argmin(test_loss_arr)]
    print(f'Best epoch: {best_epoch}')

    fig, ax = plt.subplots()
    ax.plot(epochs, np.sqrt(train_loss_arr), 'ko-', label='Train Loss')
    ax.plot(epochs, np.sqrt(test_loss_arr), 'ro-', label='Validation Loss')

    ax.legend(loc='upper right', fontsize=16)
    ax.set_xlabel('Epoch', fontsize=16)
    ax.set_ylabel('RMSE Loss', fontsize=16)

    plt.show()


if __name__ == '__main__':
    main()