n_runs.py

"""
Usage:
    run.py --hidden=<h> [options]

Options:
    --lr=<float>                        learning rate [default: 0.001]
    --bs=<int>                          batch size [default: 64]
    --seed=<int>                        random seed [default: 59]
    --epochs=<int>                      number of epochs [default: 10]
    --cuda=<int>                        use GPU id [default: 0]
    --hidden=<int>                      dimension of hidden state [default: 128]
    --layers=<int>                      layers of RNN or transformer [default: 1]
    --heads=<int>                       head number of transformer [default: 8]
    --dropout=<float>                   dropout rate [default: 0.1]
"""

import os
import random
import sys
import torch
import json
import numpy as np
import time
from docopt import docopt
from dataloader import getDataLoader
from evaluation import train_epoch, test_epoch, lossFunc


# Initialize random seeds for reproducibility
def setup_seed(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# Load question IDs from a CSV file generated by the builder
def load_question_ids(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
    q_ids = []
    for i in range(0, len(lines), 3):
        try:
            n_steps = int(lines[i].strip())
            question_ids = [int(qid) for qid in lines[i + 1].strip().split(',') if qid.strip()]
            q_ids.append(question_ids[:n_steps])
        except ValueError:
            continue
    return q_ids


# Create the model according to the selected type
def create_model(model_type, questions, hidden, layers, heads, dropout, length, device):
    if questions > 1000:
        input_dim = 16
        compressed = True
    else:
        input_dim = 2 * questions
        compressed = False

    if model_type == 'RNN':
        from RNNModel import RNNModel
        return RNNModel(input_dim, hidden, layers, questions, device), compressed
    elif model_type == 'LSTM':
        from RNNModel import LSTMModel
        return LSTMModel(input_dim, hidden, layers, questions, device), compressed
    elif model_type == 'SAKT':
        from SAKTModel import SAKTModel
        return SAKTModel(heads, length, hidden, questions, dropout), compressed
    elif model_type == 'FCN':
        from FCNModel import FCNModel
        return FCNModel(input_dim, hidden, layers, questions, device), compressed
    elif model_type == 'DKVMN':
        from DKVMNModel import DKVMNModel
        return DKVMNModel(input_dim, questions, 50, questions, device), compressed
    elif model_type == 'DKTPlus':
        from DKTPlusModel import DKTPlusModel
        return DKTPlusModel(
            input_dim, hidden, layers, questions, device,
            lambda_w1=0.5, lambda_w2=0.5, lambda_o=0.5
        ), compressed
    elif model_type == 'KQN':
        from KQNModel import KQN
        return KQN(
            n_skills=questions,
            n_hidden=hidden,
            n_rnn_hidden=hidden,
            n_mlp_hidden=hidden,
            n_rnn_layers=layers,
            rnn_type='lstm',
            device=device
        ), compressed
    else:
        raise ValueError(f"Unknown model type: {model_type}")


# Main entry point of the script
def main():
    # Parse command-line arguments
    args = docopt(__doc__)
    lr = float(args['--lr'])
    bs = int(args['--bs'])
    seed = int(args['--seed'])
    epochs = int(args['--epochs'])
    cuda = args['--cuda']
    hidden = int(args['--hidden'])
    layers = int(args['--layers'])
    heads = int(args['--heads'])
    dropout = float(args['--dropout'])

    device = torch.device(f'cuda:{cuda}' if torch.cuda.is_available() else 'cpu')

    base_dir = "dataset"
    model_types = ['RNN', 'LSTM', 'FCN', 'SAKT', 'DKVMN', 'DKTPlus', 'KQN']
    n_runs = 5  # Number of runs per model

    # Loop over each model
    for model_type in model_types:
        print(f"\n### Running experiments for model: {model_type} ###")
        model_results = []

        # Loop over each dataset
        for dataset in os.listdir(base_dir):
            dataset_path = os.path.join(base_dir, dataset)
            if not os.path.isdir(dataset_path):
                continue

            print(f"\nDataset: {dataset}")
            metrics_runs = {
                k: [] for k in
                ['auc', 'f1', 'accuracy', 'training_time', 'inference_time', 'model_size']
            }

            train_path = os.path.join(dataset_path, "builder_train.csv")
            test_path = os.path.join(dataset_path, "builder_test.csv")

            # Analyze dataset to determine the number of questions
            q_ids = load_question_ids(train_path)
            questions = max(max(seq) for seq in q_ids) + 1

            if dataset in ('statics', 'assistChall'):
                length = 500
            else:
                length = 50

            # Load DataLoaders
            trainLoader, testLoader = getDataLoader(
                train_path, test_path, bs, questions, length,
                kqn=(model_type == 'KQN')
            )

            # Repeat training for multiple runs
            for run in range(n_runs):
                print(f"\nRun {run + 1}/{n_runs}")
                setup_seed(seed + run)

                model, compressed = create_model(
                    model_type, questions, hidden, layers,
                    heads, dropout, length, device
                )

                model_size = sum(p.numel() for p in model.parameters())
                optimizer = torch.optim.Adam(model.parameters(), lr=lr)
                loss_func = lossFunc(questions, length, device, compressed=compressed)

                training_time = 0
                for epoch in range(epochs):
                    start = time.time()
                    model, optimizer = train_epoch(
                        model, trainLoader, optimizer,
                        loss_func, device, model_type
                    )
                    training_time += time.time() - start

                model.eval()
                start = time.time()
                auc, acc, f1 = test_epoch(
                    model, testLoader, loss_func, device, model_type
                )
                inference_time = time.time() - start

                metrics_runs["auc"].append(auc)
                metrics_runs["f1"].append(f1)
                metrics_runs["accuracy"].append(acc)
                metrics_runs["training_time"].append(training_time)
                metrics_runs["inference_time"].append(inference_time)
                metrics_runs["model_size"].append(model_size)

            # Compute mean and standard deviation over runs
            model_results.append({
                "dataset": dataset,
                "model": model_type,
                "avg_auc": np.mean(metrics_runs["auc"]),
                "std_auc": np.std(metrics_runs["auc"]),
                "avg_f1": np.mean(metrics_runs["f1"]),
                "std_f1": np.std(metrics_runs["f1"]),
                "avg_accuracy": np.mean(metrics_runs["accuracy"]),
                "std_accuracy": np.std(metrics_runs["accuracy"]),
                "avg_training_time": np.mean(metrics_runs["training_time"]),
                "avg_inference_time": np.mean(metrics_runs["inference_time"]),
                "model_size": np.mean(metrics_runs["model_size"])
            })

        # Save results for the current model
        os.makedirs("final_results", exist_ok=True)
        with open(f"final_results/{model_type}_summary.json", "w") as f:
            json.dump(model_results, f, indent=4)

    print("\nAll models have been trained and results have been saved.")


if __name__ == '__main__':
    main()