shunliwang

update

8bc3305 6 months ago

15.7 kB

"""
eval pretained model with multi-GPU support.
"""
import os
import numpy as np
from os.path import join
import cv2
import random
import datetime
import time
import yaml
import pickle
from tqdm import tqdm
from copy import deepcopy
from PIL import Image as pil_image
from metrics.utils import get_test_metrics
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import torch.utils.data
import torch.optim as optim
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

from dataset.abstract_dataset import DeepfakeAbstractBaseDataset
from dataset.ff_blend import FFBlendDataset
from dataset.fwa_blend import FWABlendDataset
from dataset.pair_dataset import pairDataset

from trainer.trainer import Trainer
from detectors import DETECTOR
from metrics.base_metrics_class import Recorder, calculate_acc_for_test
import metrics_retrieval.utils
from metrics_retrieval.get_metric_pro4 import *
from metrics_retrieval.get_metric import *

from collections import defaultdict

import argparse
from logger import create_logger

parser = argparse.ArgumentParser(description='Process some paths.')
parser.add_argument('--detector_path', type=str, default='/PATH/TO/resnet34.yaml', help='path to detector YAML file')
parser.add_argument("--test_dataset", nargs="+")
parser.add_argument('--weights_path', type=str, default='')
parser.add_argument('--ddp', action='store_true', help='Use DistributedDataParallel')
parser.add_argument('--use_latest', action='store_true', help='Use Latest Ckpt')
parser.add_argument('--local_rank', '--local-rank', type=int, default=-1, help='Local rank for DDP')
parser.add_argument('--test_config', type=str, default='test_config_p2.yaml', help='test_config_p2.yaml / test_config_p4.yaml')
args = parser.parse_args()


def init_seed(config, seed=None):
    if seed is None:
        if config['manualSeed'] is None:
            config['manualSeed'] = random.randint(1, 10000)
        seed = config['manualSeed']
    random.seed(seed)
    torch.manual_seed(seed)
    if config['cuda']:
        torch.cuda.manual_seed_all(seed)
    return seed


def prepare_testing_data(config, ddp=False):
    def get_test_data_loader(config, test_name):
        # update the config dictionary with the specific testing dataset
        config = config.copy()  # create a copy of config to avoid altering the original one
        config['test_dataset'] = test_name  # specify the current test dataset
        test_set = DeepfakeAbstractBaseDataset(
                config=config,
                mode='test',
            )

        # Use DistributedSampler to distribute the data
        sampler = DistributedSampler(test_set, shuffle=False) if ddp else None

        test_data_loader = \
            torch.utils.data.DataLoader(
                dataset=test_set,
                batch_size=config['test_batchSize'],
                shuffle=(sampler is None),
                num_workers=int(config['workers']),
                collate_fn=test_set.collate_fn,
                drop_last=False,
                pin_memory=True,
                sampler=sampler  # add sampler
            )
        return test_data_loader, test_set.data_dict

    test_data_loaders = {}
    test_data_dicts = {}
    for one_test_name in config['test_dataset']:
        loader, data_dict = get_test_data_loader(config, one_test_name)
        test_data_loaders[one_test_name] = loader
        test_data_dicts[one_test_name] = data_dict
    return test_data_loaders, test_data_dicts


def choose_metric(config):
    metric_scoring = config['metric_scoring']
    if metric_scoring not in ['eer', 'auc', 'acc', 'ap']:
        raise NotImplementedError('metric {} is not implemented'.format(metric_scoring))
    return metric_scoring


def test_one_dataset(model, data_loader, device, local_rank):
    # Initialize empty lists to store tensors
    prediction_lists = []
    feature_lists = []
    label_lists = []
    img_name_lists = []

    # Only the main process shows the progress bar
    pbar = tqdm(enumerate(data_loader), total=len(data_loader), disable=(local_rank != 0))

    for i, data_dict in pbar:
        # get data
        data, label, mask, landmark = data_dict['image'], data_dict['label'], data_dict['mask'], data_dict['landmark']
        img_names =[]  # data_dict['image']  # Image names are still strings and are stored separately

        # Move data to GPU (keeping the original logic)
        data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
        if mask is not None:
            data_dict['mask'] = mask.to(device)
        if landmark is not None:
            data_dict['landmark'] = landmark.to(device)

        # Model forward pass (no gradients, original logic unchanged)
        predictions = inference(model, data_dict)

        # Use append instead of extend, and concatenate later with torch.cat
        label_lists.append(data_dict['label'])  # label is a tensor, so append it directly
        prediction_lists.append(predictions['prob'])  # prob is the tensor output by the model
        feature_lists.append(predictions['feat'])  # the same applies to feat
        img_name_lists.extend(img_names)  # String lists still use extend

    # If the current process has no data (an extreme case), return empty tensors to avoid errors
    predictions_tensor = torch.cat(prediction_lists, dim=0) if prediction_lists else torch.tensor([], device=device)
    labels_tensor = torch.cat(label_lists, dim=0) if label_lists else torch.tensor([], device=device)
    feats_tensor = torch.cat(feature_lists, dim=0) if feature_lists else torch.tensor([], device=device)

    print("feats_tensor", feats_tensor.shape)

    # Return results in tensor form (image names remain a list)
    return predictions_tensor, labels_tensor, feats_tensor, img_name_lists


def test_epoch(model, test_data_loaders, test_data_dicts, device, local_rank, ddp, config, logger):
    # set model to eval mode
    model.eval()

    # define test recorder
    metrics_all_datasets = {}

    # testing for all test data
    keys = test_data_loaders.keys()
    for key in keys:

        # 1.Dataset Name
        print("Run Dataset:", key)
        # if args.local_rank == 0:
        logger.info(f"--------------- Run Dataset: {key} ---------------")
        logger.info(f"--------------- Run Dataset: {logger.log_path} ---------------")

        data_loader = test_data_loaders[key]
        data_dict = test_data_dicts[key]

        # Set the sampler epoch in DDP mode
        if ddp and hasattr(data_loader.sampler, 'set_epoch'):
            data_loader.sampler.set_epoch(0)

        # Each process computes its own portion (the return values are tensors at this point)
        predictions_tensor, labels_tensor, feats_tensor, img_names = test_one_dataset(
            model, data_loader, device, local_rank)

        # Gather results from all processes (only the main process needs the full results)
        if ddp:
            world_size = dist.get_world_size()

            # 1. Gather predictions
            all_predictions = [torch.zeros_like(predictions_tensor) for _ in range(world_size)]
            dist.all_gather(all_predictions, predictions_tensor)

            # 2. Gather labels
            all_labels = [torch.zeros_like(labels_tensor) for _ in range(world_size)]
            dist.all_gather(all_labels, labels_tensor)

            # 3. Gather features (optional)
            all_feats = [torch.zeros_like(feats_tensor) for _ in range(world_size)]
            dist.all_gather(all_feats, feats_tensor)

            all_predictions = torch.cat(all_predictions, dim=0)
            all_labels = torch.cat(all_labels, dim=0)
            all_feats = torch.cat(all_feats, dim=0)
        else:
            # In non-DDP mode, convert directly to NumPy (only once)
            all_predictions = predictions_tensor.cpu().numpy()
            all_labels = labels_tensor.cpu().numpy()
            all_feats = feats_tensor.cpu().numpy()
            #all_img_names = img_names

        # Only the main process computes metrics and outputs results
        if local_rank == 0:
            # compute metric for each dataset
            metric_one_dataset = calculate_acc_for_test(all_labels, all_predictions, config['backbone_config']['num_classes'])
            metrics_all_datasets[key] = metric_one_dataset

            # Information for each dataset
            tqdm.write(f"dataset: {key}")
            for k, v in metric_one_dataset.items():
                tqdm.write(f"{k}: {v}")
                logger.info(f"{k}: {v}")

            # save info
            pkl_save_path = os.path.join(os.path.dirname(logger.log_path), f"{key}.pkl")
            save_data = {
                "all_predictions": all_predictions.cpu().numpy(),
                "all_labels": all_labels.cpu().numpy(),
                "all_feats": all_feats.cpu().numpy(),
                "metrics": metric_one_dataset,  # Additionally save metrics for the current dataset to facilitate later analysis
                "all_names": img_names
            }
            with open(pkl_save_path, "wb") as f:
                pickle.dump(save_data, f, protocol=pickle.HIGHEST_PROTOCOL)  # Using the highest protocol is more efficient

    return metrics_all_datasets if local_rank == 0 else None


@torch.no_grad()
def inference(model, data_dict):
    from torch.cuda.amp import autocast
    with autocast(dtype=torch.float16):
        predictions = model(data_dict, inference=True)
    return predictions


def main():
    # Initialize DDP
    ddp = args.ddp
    local_rank = args.local_rank

    if ddp:
        # Initialize the process group
        torch.cuda.set_device(local_rank)
        dist.init_process_group(
            backend='nccl',
            init_method='env://',  # Read rendezvous information from environment variables (set automatically by torchrun)
            world_size=int(os.environ.get("WORLD_SIZE", 1)),  # total number of GPUs
            rank=int(os.environ.get("RANK", 0))
            )
        device = torch.device("cuda", local_rank)
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # parse options and load config
    # Model-specific configuration
    with open(args.detector_path, 'r') as f:
        config = yaml.safe_load(f)
    # Unified base configuration
    with open(f'./training/config/{args.test_config}', 'r') as f:
        config_base = yaml.safe_load(f)

    # Label dictionary shared by all datasets
    if 'label_dict' in config:
        config_base['label_dict']=config['label_dict'] # The base configuration has the highest priority
    config.update(config_base)

    weights_path = None
    # If arguments are provided, they will overwrite the yaml settings
    if args.test_dataset:
        config['test_dataset'] = args.test_dataset
    if args.weights_path:
        config['weights_path'] = args.weights_path
        weights_path = args.weights_path

    # Set the same seed for DDP
    seed = init_seed(config)
    if ddp:
        # Use a different seed offset for each process to ensure data augmentation diversity
        seed += dist.get_rank()
        init_seed(config, seed)

    # set cudnn benchmark if needed
    if config['cudnn']:
        cudnn.benchmark = True

    # Log information
    logs_test_dir = weights_path.replace("logs", "logs_test")
    # if local_rank == 0:
    # creat log
    os.makedirs(logs_test_dir, exist_ok=True)
    logger = create_logger(os.path.join(logs_test_dir, 'testing.log'))
    logger.info('Save log to {}'.format(logs_test_dir))
    # print configuration
    logger.info("--------------- Configuration ---------------")
    params_string = "Parameters: \n"
    for key, value in config.items():
        params_string += "{}: {}".format(key, value) + "\n"
    logger.info(params_string)

    # prepare the testing data loader
    test_data_loaders, test_data_dicts = prepare_testing_data(config, ddp)

    # prepare the model (detector)
    model_class = DETECTOR[config['model_name']]
    model = model_class(config).to(device)
    epoch = 0

    # Only print model parameter information on the main process
    if local_rank == 0:
        for name, param in model.named_parameters():
            print(f"{name}: {param.shape}")

    if weights_path:
        # For models containing LoRA, switch to eval mode first to avoid repeatedly stacking weights
        if 'lora' in config['model_name'].lower() or "pmoe" in config['model_name'].lower():
            model.eval()

    if weights_path:
        try:
            epoch = int(weights_path.split('/')[-1].split('.')[0].split('_')[2])
        except:
            epoch = 0

        # Automatically find the best checkpoint
        if args.use_latest:
            ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_latest.pth")
        else:
            if weights_path[-3:] == "pth":
                ckpt_path = weights_path
            else:
                ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_best.pth")
            # ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_best.pth")
        ckpt = torch.load(ckpt_path, map_location=f"cuda:{local_rank}")
        logger.info(f"Load ckpt: {ckpt_path}")

        # Remove the "module." prefix from the weights (if DDP was used during training)
        new_state_dict = {k.replace('module.', ''): v for k, v in ckpt.items()}

        model.load_state_dict(new_state_dict, strict=False)

        if local_rank == 0:
            print('===> Load checkpoint done!')
    else:
        if local_rank == 0:
            print('Fail to load the pre-trained weights')

    if ddp:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)

    # start testing
    best_metric = test_epoch(model, test_data_loaders, test_data_dicts, device, local_rank, ddp, config, logger)

    if local_rank == 0:
        print('===> Test Done!')

    # Clean up the DDP process group
    if ddp:
        dist.barrier()
        dist.destroy_process_group()

    if local_rank==0:
        #Metric test
        for prot in config['test_dataset']:
            prefix=config["weights_path"].replace("logs/", "logs_test/")
            prefix=os.path.join("/Youtu_Pangu_Security_Public_cq11/shunliwang/DeepFakeBench_DFG",prefix)
            if "protocol_2" in prot:
                pkl_file="protocol_2_test.pkl"
                RANK_MAX = 10
                seed = 42
                PKL_FILE_PATH = os.path.join(prefix, pkl_file)
                run_retrieval_evaluation(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed)
            elif "protocol_3" in prot:
                pkl_file="protocol_3_test.pkl"
                RANK_MAX = 10
                seed = 42
                PKL_FILE_PATH = os.path.join(prefix, pkl_file)
                run_retrieval_evaluation(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed)
            elif "protocol_4" in prot:
                pkl_file="protocol_4_test.pkl"
                RANK_MAX = 10
                seed = 42
                PKL_FILE_PATH = os.path.join(prefix, pkl_file)
                yaml_path="config/test_config_p4.yaml"
                run_retrieval_evaluation_p4(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed,yaml_path=yaml_path)






if __name__ == '__main__':
    main()


# 1.Useful information in the log

# 2.Create the log_test directory

# 3.Create logger text output

# 4.Save features and labels

	"""
	eval pretained model with multi-GPU support.
	"""
	import os
	import numpy as np
	from os.path import join
	import cv2
	import random
	import datetime
	import time
	import yaml
	import pickle
	from tqdm import tqdm
	from copy import deepcopy
	from PIL import Image as pil_image
	from metrics.utils import get_test_metrics
	import torch
	import torch.nn as nn
	import torch.nn.parallel
	import torch.backends.cudnn as cudnn
	import torch.nn.functional as F
	import torch.utils.data
	import torch.optim as optim
	import torch.distributed as dist
	from torch.utils.data.distributed import DistributedSampler

	from dataset.abstract_dataset import DeepfakeAbstractBaseDataset
	from dataset.ff_blend import FFBlendDataset
	from dataset.fwa_blend import FWABlendDataset
	from dataset.pair_dataset import pairDataset

	from trainer.trainer import Trainer
	from detectors import DETECTOR
	from metrics.base_metrics_class import Recorder, calculate_acc_for_test
	import metrics_retrieval.utils
	from metrics_retrieval.get_metric_pro4 import *
	from metrics_retrieval.get_metric import *

	from collections import defaultdict

	import argparse
	from logger import create_logger

	parser = argparse.ArgumentParser(description='Process some paths.')
	parser.add_argument('--detector_path', type=str, default='/PATH/TO/resnet34.yaml', help='path to detector YAML file')
	parser.add_argument("--test_dataset", nargs="+")
	parser.add_argument('--weights_path', type=str, default='')
	parser.add_argument('--ddp', action='store_true', help='Use DistributedDataParallel')
	parser.add_argument('--use_latest', action='store_true', help='Use Latest Ckpt')
	parser.add_argument('--local_rank', '--local-rank', type=int, default=-1, help='Local rank for DDP')
	parser.add_argument('--test_config', type=str, default='test_config_p2.yaml', help='test_config_p2.yaml / test_config_p4.yaml')
	args = parser.parse_args()


	def init_seed(config, seed=None):
	if seed is None:
	if config['manualSeed'] is None:
	config['manualSeed'] = random.randint(1, 10000)
	seed = config['manualSeed']
	random.seed(seed)
	torch.manual_seed(seed)
	if config['cuda']:
	torch.cuda.manual_seed_all(seed)
	return seed


	def prepare_testing_data(config, ddp=False):
	def get_test_data_loader(config, test_name):
	# update the config dictionary with the specific testing dataset
	config = config.copy() # create a copy of config to avoid altering the original one
	config['test_dataset'] = test_name # specify the current test dataset
	test_set = DeepfakeAbstractBaseDataset(
	config=config,
	mode='test',
	)

	# Use DistributedSampler to distribute the data
	sampler = DistributedSampler(test_set, shuffle=False) if ddp else None

	test_data_loader = \
	torch.utils.data.DataLoader(
	dataset=test_set,
	batch_size=config['test_batchSize'],
	shuffle=(sampler is None),
	num_workers=int(config['workers']),
	collate_fn=test_set.collate_fn,
	drop_last=False,
	pin_memory=True,
	sampler=sampler # add sampler
	)
	return test_data_loader, test_set.data_dict

	test_data_loaders = {}
	test_data_dicts = {}
	for one_test_name in config['test_dataset']:
	loader, data_dict = get_test_data_loader(config, one_test_name)
	test_data_loaders[one_test_name] = loader
	test_data_dicts[one_test_name] = data_dict
	return test_data_loaders, test_data_dicts


	def choose_metric(config):
	metric_scoring = config['metric_scoring']
	if metric_scoring not in ['eer', 'auc', 'acc', 'ap']:
	raise NotImplementedError('metric {} is not implemented'.format(metric_scoring))
	return metric_scoring


	def test_one_dataset(model, data_loader, device, local_rank):
	# Initialize empty lists to store tensors
	prediction_lists = []
	feature_lists = []
	label_lists = []
	img_name_lists = []

	# Only the main process shows the progress bar
	pbar = tqdm(enumerate(data_loader), total=len(data_loader), disable=(local_rank != 0))

	for i, data_dict in pbar:
	# get data
	data, label, mask, landmark = data_dict['image'], data_dict['label'], data_dict['mask'], data_dict['landmark']
	img_names =[] # data_dict['image'] # Image names are still strings and are stored separately

	# Move data to GPU (keeping the original logic)
	data_dict['image'], data_dict['label'] = data.to(device), label.to(device)
	if mask is not None:
	data_dict['mask'] = mask.to(device)
	if landmark is not None:
	data_dict['landmark'] = landmark.to(device)

	# Model forward pass (no gradients, original logic unchanged)
	predictions = inference(model, data_dict)

	# Use append instead of extend, and concatenate later with torch.cat
	label_lists.append(data_dict['label']) # label is a tensor, so append it directly
	prediction_lists.append(predictions['prob']) # prob is the tensor output by the model
	feature_lists.append(predictions['feat']) # the same applies to feat
	img_name_lists.extend(img_names) # String lists still use extend

	# If the current process has no data (an extreme case), return empty tensors to avoid errors
	predictions_tensor = torch.cat(prediction_lists, dim=0) if prediction_lists else torch.tensor([], device=device)
	labels_tensor = torch.cat(label_lists, dim=0) if label_lists else torch.tensor([], device=device)
	feats_tensor = torch.cat(feature_lists, dim=0) if feature_lists else torch.tensor([], device=device)

	print("feats_tensor", feats_tensor.shape)

	# Return results in tensor form (image names remain a list)
	return predictions_tensor, labels_tensor, feats_tensor, img_name_lists


	def test_epoch(model, test_data_loaders, test_data_dicts, device, local_rank, ddp, config, logger):
	# set model to eval mode
	model.eval()

	# define test recorder
	metrics_all_datasets = {}

	# testing for all test data
	keys = test_data_loaders.keys()
	for key in keys:

	# 1.Dataset Name
	print("Run Dataset:", key)
	# if args.local_rank == 0:
	logger.info(f"--------------- Run Dataset: {key} ---------------")
	logger.info(f"--------------- Run Dataset: {logger.log_path} ---------------")

	data_loader = test_data_loaders[key]
	data_dict = test_data_dicts[key]

	# Set the sampler epoch in DDP mode
	if ddp and hasattr(data_loader.sampler, 'set_epoch'):
	data_loader.sampler.set_epoch(0)

	# Each process computes its own portion (the return values are tensors at this point)
	predictions_tensor, labels_tensor, feats_tensor, img_names = test_one_dataset(
	model, data_loader, device, local_rank)

	# Gather results from all processes (only the main process needs the full results)
	if ddp:
	world_size = dist.get_world_size()

	# 1. Gather predictions
	all_predictions = [torch.zeros_like(predictions_tensor) for _ in range(world_size)]
	dist.all_gather(all_predictions, predictions_tensor)

	# 2. Gather labels
	all_labels = [torch.zeros_like(labels_tensor) for _ in range(world_size)]
	dist.all_gather(all_labels, labels_tensor)

	# 3. Gather features (optional)
	all_feats = [torch.zeros_like(feats_tensor) for _ in range(world_size)]
	dist.all_gather(all_feats, feats_tensor)

	all_predictions = torch.cat(all_predictions, dim=0)
	all_labels = torch.cat(all_labels, dim=0)
	all_feats = torch.cat(all_feats, dim=0)
	else:
	# In non-DDP mode, convert directly to NumPy (only once)
	all_predictions = predictions_tensor.cpu().numpy()
	all_labels = labels_tensor.cpu().numpy()
	all_feats = feats_tensor.cpu().numpy()
	#all_img_names = img_names

	# Only the main process computes metrics and outputs results
	if local_rank == 0:
	# compute metric for each dataset
	metric_one_dataset = calculate_acc_for_test(all_labels, all_predictions, config['backbone_config']['num_classes'])
	metrics_all_datasets[key] = metric_one_dataset

	# Information for each dataset
	tqdm.write(f"dataset: {key}")
	for k, v in metric_one_dataset.items():
	tqdm.write(f"{k}: {v}")
	logger.info(f"{k}: {v}")

	# save info
	pkl_save_path = os.path.join(os.path.dirname(logger.log_path), f"{key}.pkl")
	save_data = {
	"all_predictions": all_predictions.cpu().numpy(),
	"all_labels": all_labels.cpu().numpy(),
	"all_feats": all_feats.cpu().numpy(),
	"metrics": metric_one_dataset, # Additionally save metrics for the current dataset to facilitate later analysis
	"all_names": img_names
	}
	with open(pkl_save_path, "wb") as f:
	pickle.dump(save_data, f, protocol=pickle.HIGHEST_PROTOCOL) # Using the highest protocol is more efficient

	return metrics_all_datasets if local_rank == 0 else None


	@torch.no_grad()
	def inference(model, data_dict):
	from torch.cuda.amp import autocast
	with autocast(dtype=torch.float16):
	predictions = model(data_dict, inference=True)
	return predictions


	def main():
	# Initialize DDP
	ddp = args.ddp
	local_rank = args.local_rank

	if ddp:
	# Initialize the process group
	torch.cuda.set_device(local_rank)
	dist.init_process_group(
	backend='nccl',
	init_method='env://', # Read rendezvous information from environment variables (set automatically by torchrun)
	world_size=int(os.environ.get("WORLD_SIZE", 1)), # total number of GPUs
	rank=int(os.environ.get("RANK", 0))
	)
	device = torch.device("cuda", local_rank)
	else:
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	# parse options and load config
	# Model-specific configuration
	with open(args.detector_path, 'r') as f:
	config = yaml.safe_load(f)
	# Unified base configuration
	with open(f'./training/config/{args.test_config}', 'r') as f:
	config_base = yaml.safe_load(f)

	# Label dictionary shared by all datasets
	if 'label_dict' in config:
	config_base['label_dict']=config['label_dict'] # The base configuration has the highest priority
	config.update(config_base)

	weights_path = None
	# If arguments are provided, they will overwrite the yaml settings
	if args.test_dataset:
	config['test_dataset'] = args.test_dataset
	if args.weights_path:
	config['weights_path'] = args.weights_path
	weights_path = args.weights_path

	# Set the same seed for DDP
	seed = init_seed(config)
	if ddp:
	# Use a different seed offset for each process to ensure data augmentation diversity
	seed += dist.get_rank()
	init_seed(config, seed)

	# set cudnn benchmark if needed
	if config['cudnn']:
	cudnn.benchmark = True

	# Log information
	logs_test_dir = weights_path.replace("logs", "logs_test")
	# if local_rank == 0:
	# creat log
	os.makedirs(logs_test_dir, exist_ok=True)
	logger = create_logger(os.path.join(logs_test_dir, 'testing.log'))
	logger.info('Save log to {}'.format(logs_test_dir))
	# print configuration
	logger.info("--------------- Configuration ---------------")
	params_string = "Parameters: \n"
	for key, value in config.items():
	params_string += "{}: {}".format(key, value) + "\n"
	logger.info(params_string)

	# prepare the testing data loader
	test_data_loaders, test_data_dicts = prepare_testing_data(config, ddp)

	# prepare the model (detector)
	model_class = DETECTOR[config['model_name']]
	model = model_class(config).to(device)
	epoch = 0

	# Only print model parameter information on the main process
	if local_rank == 0:
	for name, param in model.named_parameters():
	print(f"{name}: {param.shape}")

	if weights_path:
	# For models containing LoRA, switch to eval mode first to avoid repeatedly stacking weights
	if 'lora' in config['model_name'].lower() or "pmoe" in config['model_name'].lower():
	model.eval()

	if weights_path:
	try:
	epoch = int(weights_path.split('/')[-1].split('.')[0].split('_')[2])
	except:
	epoch = 0

	# Automatically find the best checkpoint
	if args.use_latest:
	ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_latest.pth")
	else:
	if weights_path[-3:] == "pth":
	ckpt_path = weights_path
	else:
	ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_best.pth")
	# ckpt_path = os.path.join(weights_path, "test/protocol_2_test/ckpt_best.pth")
	ckpt = torch.load(ckpt_path, map_location=f"cuda:{local_rank}")
	logger.info(f"Load ckpt: {ckpt_path}")

	# Remove the "module." prefix from the weights (if DDP was used during training)
	new_state_dict = {k.replace('module.', ''): v for k, v in ckpt.items()}

	model.load_state_dict(new_state_dict, strict=False)

	if local_rank == 0:
	print('===> Load checkpoint done!')
	else:
	if local_rank == 0:
	print('Fail to load the pre-trained weights')

	if ddp:
	model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)

	# start testing
	best_metric = test_epoch(model, test_data_loaders, test_data_dicts, device, local_rank, ddp, config, logger)

	if local_rank == 0:
	print('===> Test Done!')

	# Clean up the DDP process group
	if ddp:
	dist.barrier()
	dist.destroy_process_group()

	if local_rank==0:
	#Metric test
	for prot in config['test_dataset']:
	prefix=config["weights_path"].replace("logs/", "logs_test/")
	prefix=os.path.join("/Youtu_Pangu_Security_Public_cq11/shunliwang/DeepFakeBench_DFG",prefix)
	if "protocol_2" in prot:
	pkl_file="protocol_2_test.pkl"
	RANK_MAX = 10
	seed = 42
	PKL_FILE_PATH = os.path.join(prefix, pkl_file)
	run_retrieval_evaluation(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed)
	elif "protocol_3" in prot:
	pkl_file="protocol_3_test.pkl"
	RANK_MAX = 10
	seed = 42
	PKL_FILE_PATH = os.path.join(prefix, pkl_file)
	run_retrieval_evaluation(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed)
	elif "protocol_4" in prot:
	pkl_file="protocol_4_test.pkl"
	RANK_MAX = 10
	seed = 42
	PKL_FILE_PATH = os.path.join(prefix, pkl_file)
	yaml_path="config/test_config_p4.yaml"
	run_retrieval_evaluation_p4(pkl_file_path=PKL_FILE_PATH, query_mode='10_sample_avg', rank_max=RANK_MAX,random_seed=seed,yaml_path=yaml_path)






	if __name__ == '__main__':
	main()


	# 1.Useful information in the log

	# 2.Create the log_test directory

	# 3.Create logger text output

	# 4.Save features and labels