test.py

import enum
import torch
from model import Mnist


def test(args, test_data, device):
    """
    Evaluates a trained MNIST model on the provided test dataset.

    Args:
        args: An object containing configuration parameters such as batch_size, data_dir, and save_model_path.
        test_data: The dataset to be used for testing (should be compatible with torch.utils.data.DataLoader).
        device: The device (e.g., 'cpu' or 'cuda') on which to perform computations.

    Behavior:
        - Loads the MNIST model and its trained weights from the specified path.
        - Iterates over the test dataset in batches.
        - Computes and prints the accuracy for each test batch.
        - Prints summary information about the testing process.
    """
    print(
        f"\
            ================= Testing the Minist Model =================\n \
            Batch Size: {args.batch_size}\n \
            Data Directory: {args.data_dir}\n \
            Load Model Path: {args.save_model_path}\n \
            ------------------------------------------------------------\n \
        "
    )
    test_loader = torch.utils.data.DataLoader(
        test_data, batch_size=args.batch_size, shuffle=True, num_workers=4
    )

    model = Mnist().to(device)

    model.load_state_dict(torch.load(args.save_model_path))

    model.eval()

    with torch.no_grad():
        for test_batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)

            output = model(data)

            pred = output.argmax(dim=1)
            correct = pred.eq(target).sum().item()

            print(
                f"Test Batch: {test_batch_idx}\tAccuracy: {100. * correct / len(data):.2f}%"
            )

    print("Testing completed.")