src/model_finetune.py

import pandas as pd
import numpy as np
import math
import os
from scipy.io import loadmat
from scipy.optimize import curve_fit
from sklearn.model_selection import train_test_split
import seaborn as sns
import matplotlib.pyplot as plt
from joblib import dump, load
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim.swa_utils import AveragedModel, SWALR
from torch.utils.data import DataLoader, TensorDataset
from model_regression import Mlp, MAEAndRankLoss, preprocess_data, compute_correlation_metrics, logistic_func

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if device.type == "cuda":
    torch.cuda.set_device(0)

def create_results_dataframe(data_list, network_name, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list):
    df_results = pd.DataFrame(columns=['DATASET', 'MODEL', 'SRCC', 'KRCC', 'PLCC', 'RMSE', 'SELECT_CRITERIC'])
    df_results['DATASET'] = data_list
    df_results['MODEL'] = network_name
    df_results['SRCC'] = srcc_list
    df_results['KRCC'] = krcc_list
    df_results['PLCC'] = plcc_list
    df_results['RMSE'] = rmse_list
    df_results['SELECT_CRITERIC'] = select_criteria_list
    return df_results

def process_test_set(test_data_name, metadata_path, feature_path, network_name):
    test_df = pd.read_csv(f'{metadata_path}/{test_data_name.upper()}_metadata.csv')
    if test_data_name == 'youtube_ugc' or test_data_name == 'lsvq_test' or test_data_name == 'lsvq_test_1080P':
        grey_df_test = pd.read_csv(f'{metadata_path}/greyscale_report/YOUTUBE_UGC_greyscale_metadata.csv')
        grey_indices = grey_df_test.iloc[:, 0].tolist()
        test_df = test_df.drop(index=grey_indices).reset_index(drop=True)

    test_vids = test_df['vid']
    test_scores = test_df['mos'].tolist()
    if test_data_name == 'konvid_1k' or test_data_name == 'youtube_ugc':
        test_mos_list = ((np.array(test_scores) - 1) * (99 / 4) + 1.0).tolist()
    else:
        test_mos_list = test_scores

    sorted_test_df = pd.DataFrame({
        'vid': test_df['vid'],
        'framerate': test_df['framerate'],
        'MOS': test_mos_list,
        'MOS_raw': test_scores
    })
    test_data = loadmat(f'{feature_path}/{test_data_name}_{network_name}_feats.mat')
    test_features = test_data[f'{test_data_name}']
    if test_data_name == 'youtube_ugc' or test_data_name == 'lsvq_test' or test_data_name == 'lsvq_test_1080P':
        test_features = np.delete(test_features, grey_indices, axis=0)
    print(f'num of {test_data_name} features: {len(test_features)}')
    return test_features, sorted_test_df, test_vids

def fix_state_dict(state_dict):
    new_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith('module.'):
            name = k[7:]
        elif k == 'n_averaged':
            continue
        else:
            name = k
        new_state_dict[name] = v
    return new_state_dict

def collate_to_device(batch, device):
    data, targets = zip(*batch)
    return torch.stack(data).to(device), torch.stack(targets).to(device)

def model_test(best_model, X, y, device):
    test_dataset = TensorDataset(torch.FloatTensor(X), torch.FloatTensor(y))
    test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False)

    best_model.eval()
    y_pred = []
    with torch.no_grad():
        for inputs, _ in test_loader:
            inputs = inputs.to(device)

            outputs = best_model(inputs)
            y_pred.extend(outputs.view(-1).tolist())
    return y_pred

def fine_tune_model(model, device, model_path, X_fine_tune, y_fine_tune, save_path, batch_size, epochs, loss_type, optimizer_type, initial_lr, weight_decay, use_swa, l1_w, rank_w):
    state_dict = torch.load(model_path)
    fixed_state_dict = fix_state_dict(state_dict)
    try:
        model.load_state_dict(fixed_state_dict)
    except RuntimeError as e:
        print(e)

    for param in model.parameters():
        param.requires_grad = True
    model.train().to(device) # to gpu

    fine_tune_dataset = TensorDataset(torch.FloatTensor(X_fine_tune), torch.FloatTensor(y_fine_tune))
    fine_tune_loader = DataLoader(dataset=fine_tune_dataset, batch_size=batch_size, shuffle=False)

    # initialisation of loss function, optimiser
    if loss_type == 'MAERankLoss':
        criterion = MAEAndRankLoss()
        criterion.l1_w = l1_w
        criterion.rank_w = rank_w
    else:
        nn.MSELoss()

    if optimizer_type == 'sgd':
        optimizer = optim.SGD(model.parameters(), lr=initial_lr, momentum=0.9, weight_decay=weight_decay)
        scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)# initial eta_min=1e-5
    else:
        optimizer = optim.AdamW(model.parameters(), lr=initial_lr, weight_decay=weight_decay)  # L2 Regularisation initial: 0.01, 1e-5
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.95)  # step_size=10, gamma=0.1: every 10 epochs lr*0.1
    if use_swa:
        swa_model = AveragedModel(model).to(device)
        swa_scheduler = SWALR(optimizer, swa_lr=initial_lr, anneal_strategy='cos')
    swa_start = int(epochs * 0.75) if use_swa else epochs  # SWA starts after 75% of total epochs, only set SWA start if SWA is used

    for epoch in range(epochs):
        fine_tune_loss = 0.0
        for inputs, labels in fine_tune_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels.view(-1, 1))
            loss.backward()
            optimizer.step()
            fine_tune_loss += loss.item() * inputs.size(0)

        scheduler.step()
        if use_swa and epoch >= swa_start:
            swa_model.update_parameters(model)
            swa_scheduler.step()
            print(f"Current learning rate with SWA: {swa_scheduler.get_last_lr()}")
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}, Loss: {fine_tune_loss/len(fine_tune_loader.dataset)}")

    # decide which model to evaluate: SWA model or regular model
    if use_swa and epoch >= swa_start:
        train_loader = DataLoader(dataset=fine_tune_dataset, batch_size=batch_size, shuffle=True, collate_fn=lambda x: collate_to_device(x, device))
        swa_model = swa_model.to(device)
        swa_model.eval()
        torch.optim.swa_utils.update_bn(train_loader, swa_model)
        fine_tune_model = swa_model
    else:
        fine_tune_model = model
    model_path_new = os.path.join(save_path, f"{test_data_name}_relaxvqa_{select_criteria}_fine_tuned_model.pth")
    torch.save(fine_tune_model.state_dict(), model_path_new)  # save finetuned model
    return fine_tune_model

def fine_tuned_model_test(model, device, X_test, y_test, test_data_name):
    model.train()
    state_dict = torch.load(model_path)
    fixed_state_dict = fix_state_dict(state_dict)

    model.eval()
    y_test_pred = model_test(model, X_test, y_test, device)
    y_test_pred = np.array(list(y_test_pred), dtype=float)

    if test_data_name == 'konvid_1k' or test_data_name == 'youtube_ugc':
        y_test_convert = ((np.array(y_test) - 1) / (99/4) + 1.0).tolist()
        y_test_pred_convert = ((np.array(y_test_pred) - 1) / (99/4) + 1.0).tolist()
    else:
        y_test_convert = y_test
        y_test_pred_convert = y_test_pred

    y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test_convert, y_test_pred_convert)
    test_pred_score = {'MOS': y_test_convert, 'y_test_pred': y_test_pred_convert, 'y_test_pred_logistic': y_test_pred_logistic}
    df_test_pred = pd.DataFrame(test_pred_score)
    return df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test

def wo_fine_tune_model(model, device, model_path, X_test, y_test, loss_type):
    state_dict = torch.load(model_path)
    fixed_state_dict = fix_state_dict(state_dict)
    try:
        model.load_state_dict(fixed_state_dict)
    except RuntimeError as e:
        print(e)
    model.eval().to(device) # to gpu

    if loss_type == 'MAERankLoss':
        criterion = MAEAndRankLoss()
    else:
        criterion = torch.nn.MSELoss()

    # evaluate the model
    test_dataset = TensorDataset(torch.FloatTensor(X_test), torch.FloatTensor(y_test))
    test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
    test_loss = 0.0
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, labels.view(-1, 1))
        test_loss += loss.item() * inputs.size(0)
    average_loss = test_loss / len(test_loader.dataset)
    print(f"Test Loss: {average_loss}")

    y_pred = model_test(model, X_test, y_test)
    y_pred = np.array(y_pred, dtype=float)

    if test_data_name == 'konvid_1k' or test_data_name == 'youtube_ugc':
        y_test_convert = ((np.array(y_test) - 1) / (99/4) + 1.0).tolist()
        y_test_pred_convert = ((np.array(y_pred) - 1) / (99/4) + 1.0).tolist()
    else:
        y_test_convert = y_test
        y_test_pred_convert = y_pred

    y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test_convert, y_test_pred_convert)
    test_pred_score = {'MOS': y_test_convert, 'y_test_pred': y_test_pred_convert, 'y_test_pred_logistic': y_test_pred_logistic}
    df_test_pred = pd.DataFrame(test_pred_score)
    return df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test


if __name__ == '__main__':
    # input parameters
    train_data_name = 'lsvq_train'
    network_name = 'relaxvqa'
    model_name = 'Mlp'
    select_criteria = 'byrmse'
    test_data_name = 'cvd_2014'

    metadata_path = '../metadata/'
    feature_path = f'../features/'
    save_path = '../model/'
    report_path = '../log/reported_results/'

    is_finetune = True
    data_list, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list = [], [], [], [], [], []

    # training parameters
    n_repeats = 21
    batch_size = 256
    epochs = 20
    # misc
    loss_type = 'MAERankLoss'
    optimizer_type = 'sgd'
    initial_lr = 1e-2
    weight_decay = 0.0005
    use_swa = True
    l1_w = 0.6
    rank_w = 1.0

    if is_finetune  == True:
        os.makedirs(report_path+'/fine_tune/', exist_ok=True)
        csv_name = f'{report_path}/fine_tune/{test_data_name}_{network_name}_{select_criteria}_finetune.csv'
    else:
        os.makedirs(report_path+'/fine_tune/', exist_ok=True)
        csv_name = f'{report_path}/fine_tune/{test_data_name}_{network_name}_{select_criteria}_wo_finetune.csv'
    print(f'Test dataset: {test_data_name}')
    test_features_mat, sorted_test_df, test_vids = process_test_set(test_data_name, metadata_path, feature_path, network_name)

    X_test = np.asarray(test_features_mat, dtype=float)
    y_test_data = sorted_test_df['MOS']
    with open(f"{save_path}/scaler/vid.txt", "w") as file:
        for item in sorted_test_df['vid']:
            file.write(f"{item}\n")
    with open(f"{save_path}/scaler/mos.txt", "w") as file:
        for item in y_test_data:
            file.write(f"{item}\n")
    y_test = np.array(list(y_test_data), dtype=float)
    X_test, y_test, imp, scaler = preprocess_data(X_test, y_test)
    dump(imp, f'{save_path}/scaler/{test_data_name}_imputer.pkl')
    dump(scaler, f'{save_path}/scaler/{test_data_name}_scaler.pkl')

    # get save model param
    model = Mlp(input_features=X_test.shape[1], out_features=1, drop_rate=0.2, act_layer=nn.GELU)
    model = model.to(device)
    model_path = os.path.join(save_path, f"{train_data_name}_{network_name}_{select_criteria}_trained_median_model_param_onLSVQ_TEST.pth")

    results = []
    for i in range(1, n_repeats + 1):
        print(f"{i}th repeated 80-20 hold out test")
        X_fine_tune, X_final_test, y_fine_tune, y_final_test = train_test_split(X_test, y_test, test_size=0.2, random_state=math.ceil(8.8 * i))
        if is_finetune == True:
            # test fine tuned model on the test dataset
            fine_tuned_model = fine_tune_model(model, device, model_path, X_fine_tune, y_fine_tune, save_path, batch_size,
                                               epochs, loss_type, optimizer_type, initial_lr, weight_decay, use_swa, l1_w, rank_w)

            df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = fine_tuned_model_test(fine_tuned_model, device, X_final_test, y_final_test, test_data_name)
            best_model = fine_tuned_model
        else:
            # without fine tune on the test dataset
            df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = wo_fine_tune_model(model, device, model_path, X_test, y_test, loss_type)
            print( y_test_pred_logistic)
            with open("pred_mos", "w") as file:
                for item in y_test_pred_logistic:
                    file.write(f"{item}\n")
            best_model = model

        results.append({
            'model': best_model,
            'df_pred': df_test_pred,
            'srcc': srcc_test,
            'krcc': krcc_test,
            'plcc': plcc_test,
            'rmse': rmse_test
        })
        print('\n')

    select_criteria = 'byrmse'
    if select_criteria == 'byrmse':
        sorted_results = sorted(results, key=lambda x: x['rmse'])
    median_index = len(sorted_results) // 2
    median_result = sorted_results[median_index]
    median_df_test_pred = median_result['df_pred']
    median_srcc_test = median_result['srcc']
    median_krcc_test = median_result['krcc']
    median_plcc_test = median_result['plcc']
    median_rmse_test = median_result['rmse']
    data_list.append(test_data_name)
    srcc_list.append(median_srcc_test)
    krcc_list.append(median_krcc_test)
    plcc_list.append(median_plcc_test)
    rmse_list.append(median_rmse_test)
    select_criteria_list.append(select_criteria)
    median_df_test_pred.head()

    select_criteria = 'bykrcc'
    if select_criteria == 'bykrcc':
        sorted_results = sorted(results, key=lambda x: x['krcc'])
    median_index = len(sorted_results) // 2
    median_result = sorted_results[median_index]
    median_df_test_pred = median_result['df_pred']
    median_srcc_test = median_result['srcc']
    median_krcc_test = median_result['krcc']
    median_plcc_test = median_result['plcc']
    median_rmse_test = median_result['rmse']
    data_list.append(test_data_name)
    srcc_list.append(median_srcc_test)
    krcc_list.append(median_krcc_test)
    plcc_list.append(median_plcc_test)
    rmse_list.append(median_rmse_test)
    select_criteria_list.append(select_criteria)
    median_df_test_pred.head()

    df_results = create_results_dataframe(data_list, network_name, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list)
    print(df_results.T)
    df_results.to_csv(csv_name, index=None, encoding="UTF-8")