test.py

from models.audio_models import *
from models.video_models import *
from datasets.CremaDataset import CremaDataset
from torch.utils.data import DataLoader
import warnings
import lightning as pl
from lightning.pytorch.loggers import TensorBoardLogger
import argparse
from utils.misc import Tasks
import numpy as np
import os, argparse, random
from utils.yaml_loader import *
from utils.misc import get_root_path, find_files
import re

# Ignore all kind of warnings
warnings.filterwarnings("ignore", category=Warning)

SEED = 42

def set_seed(seed):
    '''
    This method sets a unique seed for all the operations which needs the generation of (pseudo)random numbers.
    It contributes to make training and test phases predictable/deterministic, when it is possible.

    Args:
        seed (int): the seed for the generation of (pseudo)random numbers.
    '''
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    pl.seed_everything(seed, workers=True) # useful in DataLoader

set_seed(SEED)

# Command-line parsing
parser = argparse.ArgumentParser()
parser.add_argument("--conf", type=str, default=None) # optional argument to specify the configuration file
parser.add_argument("--kfold", type=int, default=None) # optional argument for k-fold cross validation
args = parser.parse_args() # read the command-line arguments

############################################################################################
# -------------------------- Main Configuration File Processing -------------------------- #
############################################################################################

# NOTE: The main configuration file is the one passed from command line

# Get the info contained into the configuration file (pt. 1)
loader = YAMLLoader(config_path=os.path.join(get_root_path(), 'conf', args.conf))
settings = loader.get_settings()
data_folder = settings.pop("data_folder")
config_path = settings.pop('config_path')
checkpoint_folder = settings.pop('checkpoint_folder')
task = settings.pop("task")

# Print the task to complete: emotion recog., gender recog. or both
if 'SingleTask' in settings['model_class']:
    task = task.upper()
    task = getattr(Tasks, task)
    print(f"Task: {task}")
else:
    print("Multitask")

# Get the info contained into the configuration file (pt. 2)
settings = loader.convert_to_obj(params=settings)
data_class = settings['data_class']
video_class = settings['video_class']
audio_class = settings['audio_class']
model_class = settings['model_class']

# Print the modality employed by the current model
if video_class and audio_class:
    print("Modality: audio-video")
elif audio_class:
    print("Modality: audio")
elif video_class:
    print("Modality: video")
else:
    print("Modality: unknown (no encoders specified)")

# Create an instance of the class associated to the model of interest
model = model_class(video=video_class, audio=audio_class, config_path=config_path, task_type=task, dataset=data_class)

###########################################################################################
# -------------------------- Sub Configuration File Processing -------------------------- #
###########################################################################################

# NOTE: The sub configuration file is the one specified by the field "config_path" of the main yml file

# Get the training info contained into the sub configuration file:
loader = YAMLLoader(config_path=config_path, key="train")
settings : dict = loader.get_settings()

# The following folder contains the files that specify the coordinates where to crop faces appearing into videos
face_folder = os.path.join(data_folder, "face_pickle") if settings["face_crop"] else None 

green_screen = settings["green_screen_test"]
print(f"Green screen test: {green_screen}")

test_set = CremaDataset(model.encoders_name, role="test", face_folder=face_folder, config_path=config_path, green_screen=green_screen, kfold=args.kfold)

batch_size = settings["batch_size"]
num_workers = settings["num_workers"]

test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=False)
# with pin_memory=True, the data will be transferred on the device (GPU) memory faster 

settings : dict = loader.get_settings(["trainer"])
settings.pop("callbacks")

name = model.name if args.kfold is None else f"{model.name}/kfold_{args.kfold}"

# Get the checkpoint corresponding to the best model version 
# (i.e. the one corresponding to the lowest validation loss value)
# NOTE we consider only the positive values for the validation loss metric
# NOTE the checkpoint files are named in the following way: "epoch=XX-val_loss=X.XX.ckpt" 
#       (in the single task scenario) or "epoch=XX-val_MT_loss=X.XX.ckpt" (in the multi task scenario)
ckpt_list = find_files(checkpoint_folder, '.ckpt')
matching_ckpts = [model for model in ckpt_list if (name in model and "loss" in model and "loss=-" not in model)]
pattern = re.compile(r"loss=([0-9]*\.?[0-9]+)")
try:
    best_ckpt = min(matching_ckpts, key=lambda path: float(pattern.search(path).group(1)))
    print(f"Best checkpoint found: {best_ckpt}")
except ValueError:
    print("Checkpoint not found!")
    exit()

tb_logger = TensorBoardLogger(checkpoint_folder, name='TEST/' + name)
trainer = pl.Trainer(logger=tb_logger, **settings)
trainer.test(model, dataloaders=test_loader, ckpt_path=best_ckpt)

#checkpoint = torch.load(best_ckpt, map_location="cpu")
#ckpt_keys = set(checkpoint['state_dict'].keys())
#model_keys = set(model.state_dict().keys())
#
## Trova le chiavi che il modello vuole ma che non sono nel checkpoint
#missing = model_keys - ckpt_keys
## Trova le chiavi che sono nel checkpoint ma che il modello non riconosce
#unexpected = ckpt_keys - model_keys
#
#print(f"MISSING: {list(missing)}")
#print(f"UNEXPECTED: {list(unexpected)}")

# ----- For backward compatibility -----
# 1. Load the checkpoint file
#checkpoint = torch.load(best_ckpt, map_location=lambda storage, loc: storage)
#
## 2. Load the state_dict into the model, ignoring the missing keys (thanks to strict=False)
## NOTE From now on, the model has already the weights loaded from the checkpoint file
#model.load_state_dict(checkpoint['state_dict'], strict=False)
#
#tb_logger = TensorBoardLogger(checkpoint_folder, name='TEST/' + name)
#trainer = pl.Trainer(logger=tb_logger, **settings)
#trainer.test(model, dataloaders=test_loader)
# --------------------------------------