ML.py

from model import MLPNet
import LoadData as LD

import os
import random
import argparse
import time  
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm
import joblib

from sklearn.metrics import accuracy_score, f1_score
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
  

import torch 
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset



def weights_init(m):
    if isinstance(m, nn.Linear):
        init.xavier_uniform_(m.weight)
        if m.bias is not None:
            init.zeros_(m.bias)

def get_device(gpu_id):
    if not torch.cuda.is_available() or gpu_id == -1:
        print("Using device: cpu"); return torch.device("cpu")

    n = torch.cuda.device_count()
    local_rank = os.getenv("LOCAL_RANK")
    if local_rank is not None and local_rank.isdigit():
        idx = int(local_rank)
    else:
        idx = gpu_id
    if idx < 0 or idx >= n:
        print(f"Rank/device {idx} out of range [0, {n-1}]; falling back to 0")
        idx = 0
    dev = torch.device(f"cuda:{idx}")
    print(f"Using device: {dev} ({torch.cuda.get_device_name(idx)})")
    return dev

def train_model(X_data, y_data, args, device, epochs=140, batch_size=128):
    model = MLPNet(num_features=X_data.shape[1],
                   num_classes=len(np.unique(y_data)),
                   dataset=args.dataset).to(device)
    model.apply(weights_init)

    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    X_tensor = torch.tensor(X_data, dtype=torch.float32)
    y_tensor = torch.tensor(y_data, dtype=torch.long)
    dataset = TensorDataset(X_tensor, y_tensor)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

    for epoch in range(epochs):
        epoch_loss = 0.0
        model.train()
        with tqdm(loader, unit="batch", desc=f"Epoch {epoch+1}/{epochs}") as tepoch:
            for data, labels in tepoch:
                data, labels = data.to(device), labels.to(device)
                optimizer.zero_grad()
                outputs = model(data)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                epoch_loss += loss.item()
                tepoch.set_postfix(loss=epoch_loss / (tepoch.n + 1))
    return model


def train_ml_models(args):
    device = get_device(args.gpu_id)
    
    X_train, _, X_clean_test, y_clean_test, clean_labels = LD.data_loader(args)
    noisy_labels_path = f'../Results/Correction/noisy-{args.dataset}-{args.exp}-{args.noise_rate}-{args.label_method}-{args.seed}.csv'
    cleansed_labels_path = f'../Results/Correction/pred-{args.dataset}-{args.exp}-{args.noise_rate}-{args.label_method}-{args.seed}.csv'
        
        
    y_train_noisy = pd.read_csv(noisy_labels_path, header=None).values.ravel()
    y_train_cleansed = pd.read_csv(cleansed_labels_path, header=None).values.ravel()

    models_dir = "../Models/ML-models"
    if not os.path.exists(models_dir):
        os.makedirs(models_dir)
    results_dir = f"../Results/ML/{args.dataset}/{args.exp}-{args.noise_rate}-{args.label_method}"
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    print("\nTraining MLPNet with noisy labels...")
    start_time = time.time()
    mlp_model_noisy = train_model(X_train, y_train_noisy, args, device, epochs=140)
    mlp_training_time_noisy = time.time() - start_time
    mlp_model_noisy.eval()
    X_clean_tensor = torch.tensor(X_clean_test, dtype=torch.float32).to(device)
    with torch.no_grad():
        outputs = mlp_model_noisy(X_clean_tensor)
        _, mlp_predictions_noisy = torch.max(outputs, 1)
    mlp_predictions_noisy_np = mlp_predictions_noisy.cpu().numpy()

    mlp_noisy_acc = accuracy_score(y_clean_test, mlp_predictions_noisy_np)
    mlp_noisy_f1 = f1_score(y_clean_test, mlp_predictions_noisy_np, average='macro', zero_division=0)

    print("\nResults for MLPNet (trained with noisy labels):")
    print(f"  Accuracy:         {mlp_noisy_acc:.4f}")
    print(f"  Macro F1 Score:   {mlp_noisy_f1:.4f}")
    print(f"  Training Time:    {mlp_training_time_noisy:.2f} seconds")

    mlp_noisy_model_path = os.path.join(models_dir, f"MLPNet_noisy_seed{args.seed}.pt")
    torch.save(mlp_model_noisy.state_dict(), mlp_noisy_model_path)

    mlp_noisy_results_file = os.path.join(results_dir, "PErealResults-MLPNet-noisy.csv")
    if not os.path.exists(mlp_noisy_results_file):
        with open(mlp_noisy_results_file, 'w') as f:
            f.write("Classifier,Dataset,LabelType,Seed,Accuracy,F1,TrainingTime\n")
    mlp_noisy_data = {
        'Classifier': 'MLPNet',
        'Dataset': args.dataset,
        'LabelType': 'noisy',
        'Seed': args.seed,
        'Accuracy': mlp_noisy_acc,
        'F1': mlp_noisy_f1,
        'TrainingTime': mlp_training_time_noisy
    }
    pd.DataFrame([mlp_noisy_data]).to_csv(mlp_noisy_results_file, mode='a', header=False, index=False)

    print("\nTraining MLPNet with cleansed labels...")
    start_time = time.time()
    mlp_model_cleansed = train_model(X_train, y_train_cleansed, args, device, epochs=140)
    mlp_training_time_cleansed = time.time() - start_time

    mlp_model_cleansed.eval()
    with torch.no_grad():
        outputs = mlp_model_cleansed(X_clean_tensor)
        _, mlp_predictions_cleansed = torch.max(outputs, 1)
    mlp_predictions_cleansed_np = mlp_predictions_cleansed.cpu().numpy()

    mlp_cleansed_acc = accuracy_score(y_clean_test, mlp_predictions_cleansed_np)
    mlp_cleansed_f1 = f1_score(y_clean_test, mlp_predictions_cleansed_np, average='macro', zero_division=0)

    print("\nResults for MLPNet (trained with cleansed labels):")
    print(f"  Accuracy:         {mlp_cleansed_acc:.4f}")
    print(f"  Macro F1 Score:   {mlp_cleansed_f1:.4f}")
    print(f"  Training Time:    {mlp_training_time_cleansed:.2f} seconds")

    mlp_cleansed_model_path = os.path.join(models_dir, f"MLPNet_cleansed_seed{args.seed}.pt")
    torch.save(mlp_model_cleansed.state_dict(), mlp_cleansed_model_path)

    mlp_cleansed_results_file = os.path.join(results_dir, "PErealResults-MLPNet-cleansed.csv")
    if not os.path.exists(mlp_cleansed_results_file):
        with open(mlp_cleansed_results_file, 'w') as f:
            f.write("Classifier,Dataset,LabelType,Seed,Accuracy,F1,TrainingTime\n")
    mlp_cleansed_data = {
        'Classifier': 'MLPNet',
        'Dataset': args.dataset,
        'LabelType': 'cleansed',
        'Seed': args.seed,
        'Accuracy': mlp_cleansed_acc,
        'F1': mlp_cleansed_f1,
        'TrainingTime': mlp_training_time_cleansed
    }
    pd.DataFrame([mlp_cleansed_data]).to_csv(mlp_cleansed_results_file, mode='a', header=False, index=False)

    
    classifiers = {
    'SVM': SVC(C=10.0, kernel='rbf', gamma='scale', random_state=args.seed),
    'RandomForest': RandomForestClassifier(
        random_state=args.seed,
        max_depth=None,          
        min_samples_split=2,     
        min_samples_leaf=1
    ),
    'DecisionTree': DecisionTreeClassifier(
        random_state=args.seed,
        max_depth=None,          
        min_samples_split=2,
        min_samples_leaf=1, 
        ccp_alpha=0.0
    ),
    'LogisticRegression': LogisticRegression(
        random_state=args.seed,
        max_iter=1000,
        C=10.0                 
    ),
    'KNeighbors': KNeighborsClassifier(
        n_neighbors=5         
    )
}

    for label_type, y_train in zip(['noisy', 'cleansed'], [y_train_noisy, y_train_cleansed]):
        print(f"\nTraining scikit-learn models with {label_type} labels...")
        for clf_name, clf in classifiers.items():
            start_time = time.time()
            clf.fit(X_train, y_train)
            training_time = time.time() - start_time

            predictions = clf.predict(X_clean_test)
            acc = accuracy_score(y_clean_test, predictions)
            f1 = f1_score(y_clean_test, predictions, average='macro', zero_division=0)

            print(f"\nResults for {clf_name} (trained with {label_type} labels):")
            print(f"  Accuracy:         {acc:.4f}")
            print(f"  Macro F1 Score:   {f1:.4f}")
            print(f"  Training Time:    {training_time:.2f} seconds")

            result_data = {
                'Classifier': clf_name,
                'Dataset': args.dataset,
                'LabelType': label_type,
                'Seed': args.seed,
                'Accuracy': acc,
                'F1': f1,
                'TrainingTime': training_time
            }
            clf_results_file = os.path.join(results_dir, f"PErealResults-{clf_name}-{label_type}.csv")
            if not os.path.exists(clf_results_file):
                with open(clf_results_file, 'w') as f:
                    f.write("Classifier,Dataset,LabelType,Seed,Accuracy,F1,TrainingTime\n")
            pd.DataFrame([result_data]).to_csv(clf_results_file, mode='a', header=False, index=False)

            model_filename = os.path.join(models_dir, f"{clf_name}_{label_type}_seed{args.seed}.joblib")
            joblib.dump(clf, model_filename)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Train ML models (MLP and lightweight classifiers) on cleansed and noisy labels")
    parser.add_argument('--seed', type=int, default=1, help="Random seed")
    parser.add_argument('--dataset', type=str, default='windows_pe', help='Dataset name')
    parser.add_argument('--exp', type=str, default='random', help='Experiment name')
    parser.add_argument('--noise_rate', type=float, default=0.0, help='Noise Rate')
    parser.add_argument('--label_method', type=str, default='Label', help='Label Method: Label or y_noisy')
    parser.add_argument('--gpu_id', type=int, default=0) # gpu id
    return parser.parse_args()


def main():
    args = parse_args()
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(args.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    
    train_ml_models(args)


if __name__ == '__main__':
    main()