DAGsHub-Tutorial-MNIST/code/eval.py

"""
Evaluate model performance
"""
import pickle
import json
import numpy as np
from sklearn.metrics import accuracy_score
import torch


def eval_model():
    # Load test data
    print("Loading data and model...")
    test_data = np.load('./data/processed_test_data.npy')

    # Load trained model
    with open('./data/model.pkl', 'rb') as f:
        model = pickle.load(f)

    # Switch model to evaluation (inference) mode
    model.eval()

    print("done.")

    # Divide loaded data-set into data and labels
    labels = test_data[:, 0]
    data = torch.Tensor(test_data[:, 1:].reshape([test_data.shape[0], 1, 28, 28]))

    # Run model on test data
    print("Running model on test data...")
    predictions = model(data).max(1, keepdim=True)[1].cpu().data.numpy()
    print("done.")

    # Calculate metric scores
    print("Calculating metrics...")
    metrics = {'accuracy': accuracy_score(labels, predictions)}

    # Save metrics to json file
    with open('./metrics/eval.json', 'w') as f:
        json.dump(metrics, f)
    print("done.")


if __name__ == '__main__':
    eval_model()