Upload 42 files

config/__init__.py

@@ -0,0 +1,7 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)

config/constants.py

@@ -0,0 +1,15 @@
+import os
+
+# Path to the directory containing the constants.py file
+CONFIGS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# The base directory for UCF-related files, i.e., UCF directory
+UCF_BASE_PATH = os.path.abspath(os.path.join(CONFIGS_DIR, ".."))  # Points to bitmind-subnet/base_miner/UCF/
+# Absolute paths for the required files and directories
+CONFIG_PATH = os.path.join(CONFIGS_DIR, "ucf.yaml")  # Path to the ucf.yaml file
+WEIGHTS_DIR = os.path.join(UCF_BASE_PATH, "weights/") # Path to pretrained weights directory
+
+HF_REPO = "bitmind/ucf"
+BACKBONE_CKPT = "xception_best.pth"
+
+DLIB_FACE_PREDICTOR_PATH = os.path.abspath(os.path.join(UCF_BASE_PATH, "../../utils/dlib_tools/shape_predictor_81_face_landmarks.dat"))

config/pretrained_config.yaml

@@ -0,0 +1,94 @@
+SWA: false
+backbone_config:
+  dropout: false
+  inc: 3
+  mode: adjust_channel
+  num_classes: 2
+backbone_name: xception
+compression: c23
+cuda: true
+cudnn: true
+dataset_json_folder: preprocessing/dataset_json_v3
+dataset_meta:
+  fake:
+  - create_splits: false
+    path: bitmind/celeb-a-hq___stable-diffusion-xl-base-1.0___256_training_faces
+  - create_splits: false
+    path: bitmind/ffhq-256___stable-diffusion-xl-base-1.0_training_faces
+  real:
+  - create_splits: false
+    path: bitmind/celeb-a-hq_training_faces
+  - create_splits: false
+    path: bitmind/ffhq-256_training_faces
+ddp: false
+dry_run: false
+encoder_feat_dim: 512
+faces_only: true
+frame_num:
+  test: 32
+  train: 32
+lmdb: true
+lmdb_dir: ./datasets/lmdb
+local_rank: 0
+log_dir: ./logs/training/ucf_2024-09-17-16-44-50
+logdir: ./logs
+loss_func:
+  cls_loss: cross_entropy
+  con_loss: contrastive_regularization
+  rec_loss: l1loss
+  spe_loss: cross_entropy
+losstype: null
+lr_scheduler: null
+manualSeed: 1024
+mean:
+- 0.5
+- 0.5
+- 0.5
+metric_scoring: auc
+mode: train
+model_name: ucf
+nEpochs: 2
+optimizer:
+  adam:
+    amsgrad: false
+    beta1: 0.9
+    beta2: 0.999
+    eps: 1.0e-08
+    lr: 0.0002
+    weight_decay: 0.0005
+  sgd:
+    lr: 0.0002
+    momentum: 0.9
+    weight_decay: 0.0005
+  type: adam
+pretrained: ../weights/xception_best.pth
+rec_iter: 100
+resolution: 256
+rgb_dir: ./datasets/rgb
+save_avg: true
+save_ckpt: true
+save_epoch: 1
+save_feat: true
+specific_task_number: 2
+split_transforms:
+  test:
+    name: base_transforms
+  train:
+    name: random_aug_transforms
+  validation:
+    name: base_transforms
+start_epoch: 0
+std:
+- 0.5
+- 0.5
+- 0.5
+test_batchSize: 32
+train_batchSize: 32
+train_dataset:
+- bitmind/celeb-a-hq_training_faces
+- bitmind/ffhq-256_training_faces
+- bitmind/celeb-a-hq___stable-diffusion-xl-base-1.0___256_training_faces
+- bitmind/ffhq-256___stable-diffusion-xl-base-1.0_training_faces
+with_landmark: false
+with_mask: false
+workers: 7

config/pretrained_face_config.yaml

@@ -0,0 +1,94 @@
+SWA: false
+backbone_config:
+  dropout: false
+  inc: 3
+  mode: adjust_channel
+  num_classes: 2
+backbone_name: xception
+compression: c23
+cuda: true
+cudnn: true
+dataset_json_folder: preprocessing/dataset_json_v3
+dataset_meta:
+  fake:
+  - create_splits: false
+    path: bitmind/celeb-a-hq___stable-diffusion-xl-base-1.0___256_training_faces
+  - create_splits: false
+    path: bitmind/ffhq-256___stable-diffusion-xl-base-1.0_training_faces
+  real:
+  - create_splits: false
+    path: bitmind/celeb-a-hq_training_faces
+  - create_splits: false
+    path: bitmind/ffhq-256_training_faces
+ddp: false
+dry_run: false
+encoder_feat_dim: 512
+faces_only: true
+frame_num:
+  test: 32
+  train: 32
+lmdb: true
+lmdb_dir: ./datasets/lmdb
+local_rank: 0
+log_dir: ./logs/training/ucf_2024-09-17-16-44-50
+logdir: ./logs
+loss_func:
+  cls_loss: cross_entropy
+  con_loss: contrastive_regularization
+  rec_loss: l1loss
+  spe_loss: cross_entropy
+losstype: null
+lr_scheduler: null
+manualSeed: 1024
+mean:
+- 0.5
+- 0.5
+- 0.5
+metric_scoring: auc
+mode: train
+model_name: ucf
+nEpochs: 2
+optimizer:
+  adam:
+    amsgrad: false
+    beta1: 0.9
+    beta2: 0.999
+    eps: 1.0e-08
+    lr: 0.0002
+    weight_decay: 0.0005
+  sgd:
+    lr: 0.0002
+    momentum: 0.9
+    weight_decay: 0.0005
+  type: adam
+pretrained: ../weights/xception_best.pth
+rec_iter: 100
+resolution: 256
+rgb_dir: ./datasets/rgb
+save_avg: true
+save_ckpt: true
+save_epoch: 1
+save_feat: true
+specific_task_number: 2
+split_transforms:
+  test:
+    name: base_transforms
+  train:
+    name: random_aug_transforms
+  validation:
+    name: base_transforms
+start_epoch: 0
+std:
+- 0.5
+- 0.5
+- 0.5
+test_batchSize: 32
+train_batchSize: 32
+train_dataset:
+- bitmind/celeb-a-hq_training_faces
+- bitmind/ffhq-256_training_faces
+- bitmind/celeb-a-hq___stable-diffusion-xl-base-1.0___256_training_faces
+- bitmind/ffhq-256___stable-diffusion-xl-base-1.0_training_faces
+with_landmark: false
+with_mask: false
+workers: 7

config/train_config.yaml

@@ -0,0 +1,9 @@
+mode: train
+lmdb: True
+dry_run: false
+rgb_dir: './datasets/rgb'
+lmdb_dir:  './datasets/lmdb'
+dataset_json_folder: './preprocessing/dataset_json'
+SWA: False
+save_avg: True
+log_dir: ./logs/training/

config/ucf.yaml

@@ -0,0 +1,73 @@
+# log dir 
+log_dir: ../debug_logs/ucf
+
+# model setting
+pretrained: ../weights/xception_best.pth   # path to a pre-trained model, if using one
+model_name: ucf   # model name
+backbone_name: xception  # backbone name
+encoder_feat_dim: 512  # feature dimension of the backbone
+  
+#backbone setting
+backbone_config:
+  mode: adjust_channel
+  num_classes: 2
+  inc: 3
+  dropout: false
+
+compression: c23  # compression-level for videos
+train_batchSize: 32   # training batch size
+test_batchSize: 32   # test batch size
+workers: 8   # number of data loading workers
+frame_num: {'train': 32, 'test': 32}   # number of frames to use per video in training and testing
+resolution: 256   # resolution of output image to network
+with_mask: false   # whether to include mask information in the input
+with_landmark: false   # whether to include facial landmark information in the input
+save_ckpt: true   # whether to save checkpoint
+save_feat: true   # whether to save features
+specific_task_number: 2 # default num datasets in FF++ used by DFB, overwritten in training
+
+# mean and std for normalization
+mean: [0.5, 0.5, 0.5]
+std: [0.5, 0.5, 0.5]
+
+# optimizer config
+optimizer:
+  # choose between 'adam' and 'sgd'
+  type: adam
+  adam:
+    lr: 0.0002  # learning rate
+    beta1: 0.9  # beta1 for Adam optimizer
+    beta2: 0.999 # beta2 for Adam optimizer
+    eps: 0.00000001  # epsilon for Adam optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+    amsgrad: false
+  sgd:
+    lr: 0.0002  # learning rate
+    momentum: 0.9  # momentum for SGD optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+
+# training config
+lr_scheduler: null   # learning rate scheduler
+nEpochs: 20   # number of epochs to train for
+start_epoch: 0   # manual epoch number (useful for restarts)
+save_epoch: 1   # interval epochs for saving models
+rec_iter: 100   # interval iterations for recording
+logdir: ./logs   # folder to output images and logs
+manualSeed: 1024   # manual seed for random number generation
+save_ckpt: false   # whether to save checkpoint
+
+# loss function
+loss_func:
+ cls_loss: cross_entropy   # loss function to use
+ spe_loss: cross_entropy
+ con_loss: contrastive_regularization
+ rec_loss: l1loss
+losstype: null
+
+# metric
+metric_scoring: auc   # metric for evaluation (auc, acc, eer, ap)
+
+# cuda
+
+cuda: true   # whether to use CUDA acceleration
+cudnn: true   # whether to use CuDNN for convolution operations

config/xception.yaml

@@ -0,0 +1,86 @@
+# log dir 
+log_dir: /data/home/zhiyuanyan/DeepfakeBench/logs/testing_bench
+
+# model setting
+pretrained: /data/home/zhiyuanyan/DeepfakeBench/training/pretrained/xception-b5690688.pth   # path to a pre-trained model, if using one
+model_name: xception   # model name
+backbone_name: xception  # backbone name
+
+#backbone setting
+backbone_config:
+  mode: original
+  num_classes: 2
+  inc: 3
+  dropout: false
+
+# dataset
+all_dataset: [FaceForensics++, FF-F2F, FF-DF, FF-FS, FF-NT, FaceShifter, DeepFakeDetection, Celeb-DF-v1, Celeb-DF-v2, DFDCP, DFDC, DeeperForensics-1.0, UADFV]
+train_dataset: [FaceForensics++]
+test_dataset: [FaceForensics++, DeepFakeDetection]
+
+compression: c23  # compression-level for videos
+train_batchSize: 32   # training batch size
+test_batchSize: 32   # test batch size
+workers: 8   # number of data loading workers
+frame_num: {'train': 32, 'test': 32}   # number of frames to use per video in training and testing
+resolution: 256   # resolution of output image to network
+with_mask: false   # whether to include mask information in the input
+with_landmark: false   # whether to include facial landmark information in the input
+
+
+# data augmentation
+use_data_augmentation: true  # Add this flag to enable/disable data augmentation
+data_aug:
+  flip_prob: 0.5
+  rotate_prob: 0.0
+  rotate_limit: [-10, 10]
+  blur_prob: 0.5
+  blur_limit: [3, 7]
+  brightness_prob: 0.5
+  brightness_limit: [-0.1, 0.1]
+  contrast_limit: [-0.1, 0.1]
+  quality_lower: 40
+  quality_upper: 100
+
+# mean and std for normalization
+mean: [0.5, 0.5, 0.5]
+std: [0.5, 0.5, 0.5]
+
+# optimizer config
+optimizer:
+  # choose between 'adam' and 'sgd'
+  type: adam
+  adam:
+    lr: 0.0002  # learning rate
+    beta1: 0.9  # beta1 for Adam optimizer
+    beta2: 0.999 # beta2 for Adam optimizer
+    eps: 0.00000001  # epsilon for Adam optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+    amsgrad: false
+  sgd:
+    lr: 0.0002  # learning rate
+    momentum: 0.9  # momentum for SGD optimizer
+    weight_decay: 0.0005  # weight decay for regularization
+
+# training config
+lr_scheduler: null   # learning rate scheduler
+nEpochs: 10   # number of epochs to train for
+start_epoch: 0   # manual epoch number (useful for restarts)
+save_epoch: 1   # interval epochs for saving models
+rec_iter: 100   # interval iterations for recording
+logdir: ./logs   # folder to output images and logs
+manualSeed: 1024   # manual seed for random number generation
+save_ckpt: true   # whether to save checkpoint
+save_feat: true   # whether to save features
+
+# loss function
+loss_func: cross_entropy   # loss function to use
+losstype: null
+
+# metric
+metric_scoring: auc   # metric for evaluation (auc, acc, eer, ap)
+
+# cuda
+
+cuda: true   # whether to use CUDA acceleration
+cudnn: true   # whether to use CuDNN for convolution operations

detectors/__init__.py

@@ -0,0 +1,11 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)
+
+from metrics.registry import DETECTOR
+
+from .ucf_detector import UCFDetector

detectors/base_detector.py

@@ -0,0 +1,71 @@
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+# description: Abstract Class for the Deepfake Detector
+
+import abc
+import torch
+import torch.nn as nn
+from typing import Union
+
+class AbstractDetector(nn.Module, metaclass=abc.ABCMeta):
+    """
+    All deepfake detectors should subclass this class.
+    """
+    def __init__(self, config=None, load_param: Union[bool, str] = False):
+        """
+        config:   (dict)
+            configurations for the model
+        load_param:  (False | True | Path(str))
+            False Do not read; True Read the default path; Path Read the required path
+        """
+        super().__init__()
+
+    @abc.abstractmethod
+    def features(self, data_dict: dict) -> torch.tensor:
+        """
+        Returns the features from the backbone given the input data.
+        """
+        pass
+
+    @abc.abstractmethod
+    def forward(self, data_dict: dict, inference=False) -> dict:
+        """
+        Forward pass through the model, returning the prediction dictionary.
+        """
+        pass
+
+    @abc.abstractmethod
+    def classifier(self, features: torch.tensor) -> torch.tensor:
+        """
+        Classifies the features into classes.
+        """
+        pass
+
+    @abc.abstractmethod
+    def build_backbone(self, config):
+        """
+        Builds the backbone of the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def build_loss(self, config):
+        """
+        Builds the loss function for the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+        """
+        Returns the losses for the model.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
+        """
+        Returns the training metrics for the model.
+        """
+        pass

detectors/ucf_detector.py

@@ -0,0 +1,472 @@
+'''
+# Source: https://github.com/SCLBD/DeepfakeBench/blob/main/training/detectors/ucf_detector.py
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+# description: Class for the UCFDetector
+
+Functions in the Class are summarized as:
+1. __init__: Initialization
+2. build_backbone: Backbone-building
+3. build_loss: Loss-function-building
+4. features: Feature-extraction
+5. classifier: Classification
+6. get_losses: Loss-computation
+7. get_train_metrics: Training-metrics-computation
+8. get_test_metrics: Testing-metrics-computation
+9. forward: Forward-propagation
+
+Reference:
+@article{yan2023ucf,
+  title={UCF: Uncovering Common Features for Generalizable Deepfake Detection},
+  author={Yan, Zhiyuan and Zhang, Yong and Fan, Yanbo and Wu, Baoyuan},
+  journal={arXiv preprint arXiv:2304.13949},
+  year={2023}
+}
+'''
+
+import os
+import datetime
+import logging
+import random
+import numpy as np
+from sklearn import metrics
+from typing import Union
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.nn import DataParallel
+from torch.utils.tensorboard import SummaryWriter
+
+from metrics.base_metrics_class import calculate_metrics_for_train
+
+from .base_detector import AbstractDetector
+from arena.detectors.UCF.detectors import DETECTOR
+from networks import BACKBONE
+from loss import LOSSFUNC
+
+logger = logging.getLogger(__name__)
+
+@DETECTOR.register_module(module_name='ucf')
+class UCFDetector(AbstractDetector):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_classes = config['backbone_config']['num_classes']
+        self.encoder_feat_dim = config['encoder_feat_dim']
+        self.half_fingerprint_dim = self.encoder_feat_dim//2
+
+        self.encoder_f = self.build_backbone(config)
+        self.encoder_c = self.build_backbone(config)
+
+        self.loss_func = self.build_loss(config)
+        self.prob, self.label = [], []
+        self.correct, self.total = 0, 0
+        
+        # basic function
+        self.lr = nn.LeakyReLU(inplace=True)
+        self.do = nn.Dropout(0.2)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+
+        # conditional gan
+        self.con_gan = Conditional_UNet()
+
+        # head
+        specific_task_number = config['specific_task_number']
+
+        self.head_spe = Head(
+            in_f=self.half_fingerprint_dim, 
+            hidden_dim=self.encoder_feat_dim,
+            out_f=specific_task_number
+        )
+        self.head_sha = Head(
+            in_f=self.half_fingerprint_dim,
+            hidden_dim=self.encoder_feat_dim, 
+            out_f=self.num_classes
+        )
+        self.block_spe = Conv2d1x1(
+            in_f=self.encoder_feat_dim,
+            hidden_dim=self.half_fingerprint_dim, 
+            out_f=self.half_fingerprint_dim
+        )
+        self.block_sha = Conv2d1x1(
+            in_f=self.encoder_feat_dim, 
+            hidden_dim=self.half_fingerprint_dim, 
+            out_f=self.half_fingerprint_dim
+        )
+        
+    def build_backbone(self, config):
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        pretrained_path = os.path.join(current_dir, config['pretrained'])
+        # prepare the backbone
+        backbone_class = BACKBONE[config['backbone_name']]
+        model_config = config['backbone_config']
+        backbone = backbone_class(model_config)
+        # if donot load the pretrained weights, fail to get good results
+        state_dict = torch.load(pretrained_path)
+        for name, weights in state_dict.items():
+            if 'pointwise' in name:
+                state_dict[name] = weights.unsqueeze(-1).unsqueeze(-1)
+        state_dict = {k:v for k, v in state_dict.items() if 'fc' not in k}
+        backbone.load_state_dict(state_dict, False)
+        logger.info('Load pretrained model successfully!')
+        return backbone
+    
+    def build_loss(self, config):
+        cls_loss_class = LOSSFUNC[config['loss_func']['cls_loss']]
+        spe_loss_class = LOSSFUNC[config['loss_func']['spe_loss']]
+        con_loss_class = LOSSFUNC[config['loss_func']['con_loss']]
+        rec_loss_class = LOSSFUNC[config['loss_func']['rec_loss']]
+        cls_loss_func = cls_loss_class()
+        spe_loss_func = spe_loss_class()
+        con_loss_func = con_loss_class(margin=3.0)
+        rec_loss_func = rec_loss_class()
+        loss_func = {
+            'cls': cls_loss_func, 
+            'spe': spe_loss_func,
+            'con': con_loss_func,
+            'rec': rec_loss_func,
+        }
+        return loss_func
+    
+    def features(self, data_dict: dict) -> torch.tensor:
+        cat_data = data_dict['image']
+        # encoder
+        f_all = self.encoder_f.features(cat_data)
+        c_all = self.encoder_c.features(cat_data)
+        feat_dict = {'forgery': f_all, 'content': c_all}
+        return feat_dict
+
+    def classifier(self, features: torch.tensor) -> torch.tensor:
+        # classification, multi-task
+        # split the features into the specific and common forgery
+        f_spe = self.block_spe(features)
+        f_share = self.block_sha(features)
+        return f_spe, f_share
+    
+    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+        if 'label_spe' in data_dict and 'recontruction_imgs' in pred_dict:
+            return self.get_train_losses(data_dict, pred_dict)
+        else:  # test mode
+            return self.get_test_losses(data_dict, pred_dict)
+
+    def get_train_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+        # get combined, real, fake imgs
+        cat_data = data_dict['image']
+        real_img, fake_img = cat_data.chunk(2, dim=0)
+        # get the reconstruction imgs
+        reconstruction_image_1, \
+        reconstruction_image_2, \
+        self_reconstruction_image_1, \
+        self_reconstruction_image_2 \
+            = pred_dict['recontruction_imgs']
+        # get label
+        label = data_dict['label']
+        label_spe = data_dict['label_spe']
+        # get pred
+        pred = pred_dict['cls']
+        pred_spe = pred_dict['cls_spe']
+
+        # 1. classification loss for common features
+        loss_sha = self.loss_func['cls'](pred, label)
+        
+        # 2. classification loss for specific features
+        loss_spe = self.loss_func['spe'](pred_spe, label_spe)
+
+        # 3. reconstruction loss
+        self_loss_reconstruction_1 = self.loss_func['rec'](fake_img, self_reconstruction_image_1)
+        self_loss_reconstruction_2 = self.loss_func['rec'](real_img, self_reconstruction_image_2)
+        cross_loss_reconstruction_1 = self.loss_func['rec'](fake_img, reconstruction_image_2)
+        cross_loss_reconstruction_2 = self.loss_func['rec'](real_img, reconstruction_image_1)
+        loss_reconstruction = \
+            self_loss_reconstruction_1 + self_loss_reconstruction_2 + \
+            cross_loss_reconstruction_1 + cross_loss_reconstruction_2
+
+        # 4. constrative loss
+        common_features = pred_dict['feat']
+        specific_features = pred_dict['feat_spe']
+        loss_con = self.loss_func['con'](common_features, specific_features, label_spe)
+
+        # 5. total loss
+        loss = loss_sha + 0.1*loss_spe + 0.3*loss_reconstruction + 0.05*loss_con
+        loss_dict = {
+            'overall': loss,
+            'common': loss_sha,
+            'specific': loss_spe,
+            'reconstruction': loss_reconstruction,
+            'contrastive': loss_con,
+        }
+        return loss_dict
+
+    def get_test_losses(self, data_dict: dict, pred_dict: dict) -> dict:
+        # get label
+        label = data_dict['label']
+        # get pred
+        pred = pred_dict['cls']
+        # for test mode, only classification loss for common features
+        loss = self.loss_func['cls'](pred, label)
+        loss_dict = {'common': loss}
+        return loss_dict
+
+    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
+        def get_accracy(label, output):
+            _, prediction = torch.max(output, 1)    # argmax
+            correct = (prediction == label).sum().item()
+            accuracy = correct / prediction.size(0)
+            return accuracy
+        
+        # get pred and label
+        label = data_dict['label']
+        pred = pred_dict['cls']
+        label_spe = data_dict['label_spe']
+        pred_spe = pred_dict['cls_spe']
+
+        # compute metrics for batch data
+        auc, eer, acc, ap = calculate_metrics_for_train(label.detach(), pred.detach())
+        acc_spe = get_accracy(label_spe.detach(), pred_spe.detach())
+        metric_batch_dict = {'acc': acc, 'acc_spe': acc_spe, 'auc': auc, 'eer': eer, 'ap': ap}
+        # we dont compute the video-level metrics for training
+        return metric_batch_dict
+
+    def forward(self, data_dict: dict, inference=False) -> dict:
+        # split the features into the content and forgery
+        features = self.features(data_dict)
+        forgery_features, content_features = features['forgery'], features['content']
+        # get the prediction by classifier (split the common and specific forgery)
+        f_spe, f_share = self.classifier(forgery_features)
+
+        if inference:
+            # inference only consider share loss
+            out_sha, sha_feat = self.head_sha(f_share)
+            out_spe, spe_feat = self.head_spe(f_spe)
+            prob_sha = torch.softmax(out_sha, dim=1)[:, 1]
+            self.prob.append(
+                prob_sha
+                .detach()
+                .squeeze()
+                .cpu()
+                .numpy()
+            )
+            _, prediction_class = torch.max(out_sha, 1)
+            if 'label' in data_dict: 
+                self.label.append(
+                    data_dict['label']
+                    .detach()
+                    .squeeze()
+                    .cpu()
+                    .numpy()
+                )
+                # deal with acc
+                common_label = (data_dict['label'] >= 1)
+                correct = (prediction_class == common_label).sum().item()
+                self.correct += correct
+                self.total += data_dict['label'].size(0)
+
+            pred_dict = {'cls': out_sha, 'feat': sha_feat}
+            return  pred_dict
+
+        bs = f_share.size(0)
+        # using idx aug in the training mode
+        aug_idx = random.random()
+        if aug_idx < 0.7:
+            # real
+            idx_list = list(range(0, bs//2))
+            random.shuffle(idx_list)
+            f_share[0: bs//2] = f_share[idx_list]
+            # fake
+            idx_list = list(range(bs//2, bs))
+            random.shuffle(idx_list)
+            f_share[bs//2: bs] = f_share[idx_list]
+        
+        # concat spe and share to obtain new_f_all
+        f_all = torch.cat((f_spe, f_share), dim=1)
+        
+        # reconstruction loss
+        f2, f1 = f_all.chunk(2, dim=0)
+        c2, c1 = content_features.chunk(2, dim=0)
+
+        # ==== self reconstruction ==== #
+        # f1 + c1 -> f11, f11 + c1 -> near~I1
+        self_reconstruction_image_1 = self.con_gan(f1, c1)
+
+        # f2 + c2 -> f2, f2 + c2 -> near~I2
+        self_reconstruction_image_2 = self.con_gan(f2, c2)
+
+        # ==== cross combine ==== #
+        reconstruction_image_1 = self.con_gan(f1, c2)
+        reconstruction_image_2 = self.con_gan(f2, c1)
+
+        # head for spe and sha
+        out_spe, spe_feat = self.head_spe(f_spe)
+        out_sha, sha_feat = self.head_sha(f_share)
+
+        # get the probability of the pred
+        prob_sha = torch.softmax(out_sha, dim=1)[:, 1]
+        prob_spe = torch.softmax(out_spe, dim=1)[:, 1]
+
+        # build the prediction dict for each output
+        pred_dict = {
+            'cls': out_sha, 
+            'prob': prob_sha, 
+            'feat': sha_feat,
+            'cls_spe': out_spe,
+            'prob_spe': prob_spe,
+            'feat_spe': spe_feat,
+            'feat_content': content_features,
+            'recontruction_imgs': (
+                reconstruction_image_1,
+                reconstruction_image_2,
+                self_reconstruction_image_1, 
+                self_reconstruction_image_2
+            )
+        }
+        return pred_dict
+
+def sn_double_conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.utils.spectral_norm(
+            nn.Conv2d(in_channels, in_channels, 3, padding=1)),
+        nn.utils.spectral_norm(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, stride=2)),
+        nn.LeakyReLU(0.2, inplace=True)
+    )
+
+def r_double_conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(out_channels, out_channels, 3, padding=1),
+        nn.ReLU(inplace=True)
+    )
+
+class AdaIN(nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        # self.l1 = nn.Linear(num_classes, in_channel*4, bias=True) #bias is good :)
+
+    def c_norm(self, x, bs, ch, eps=1e-7):
+        # assert isinstance(x, torch.cuda.FloatTensor)
+        x_var = x.var(dim=-1) + eps
+        x_std = x_var.sqrt().view(bs, ch, 1, 1)
+        x_mean = x.mean(dim=-1).view(bs, ch, 1, 1)
+        return x_std, x_mean
+
+    def forward(self, x, y):
+        assert x.size(0)==y.size(0)
+        size = x.size()
+        bs, ch = size[:2]
+        x_ = x.view(bs, ch, -1)
+        y_ = y.reshape(bs, ch, -1)
+        x_std, x_mean = self.c_norm(x_, bs, ch, eps=self.eps)
+        y_std, y_mean = self.c_norm(y_, bs, ch, eps=self.eps)
+        out =   ((x - x_mean.expand(size)) / x_std.expand(size)) \
+                * y_std.expand(size) + y_mean.expand(size)
+        return out
+
+class Conditional_UNet(nn.Module):
+
+    def init_weight(self, std=0.2):
+        for m in self.modules():
+            cn = m.__class__.__name__
+            if cn.find('Conv') != -1:
+                m.weight.data.normal_(0., std)
+            elif cn.find('Linear') != -1:
+                m.weight.data.normal_(1., std)
+                m.bias.data.fill_(0)
+
+    def __init__(self):
+        super(Conditional_UNet, self).__init__()
+
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.maxpool = nn.MaxPool2d(2)
+        self.dropout = nn.Dropout(p=0.3)
+        #self.dropout_half = HalfDropout(p=0.3)
+        
+        self.adain3 = AdaIN()
+        self.adain2 = AdaIN()
+        self.adain1 = AdaIN()
+
+        self.dconv_up3 = r_double_conv(512, 256)
+        self.dconv_up2 = r_double_conv(256, 128)
+        self.dconv_up1 = r_double_conv(128, 64)
+        
+        self.conv_last = nn.Conv2d(64, 3, 1)
+        self.up_last = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True)
+        self.activation = nn.Tanh()
+        #self.init_weight() 
+        
+    def forward(self, c, x):  # c is the style and x is the content
+        x = self.adain3(x, c)
+        x = self.upsample(x)
+        x = self.dropout(x)
+        x = self.dconv_up3(x)
+        c = self.upsample(c)
+        c = self.dropout(c)
+        c = self.dconv_up3(c)
+
+        x = self.adain2(x, c)
+        x = self.upsample(x)        
+        x = self.dropout(x)     
+        x = self.dconv_up2(x)
+        c = self.upsample(c)        
+        c = self.dropout(c)     
+        c = self.dconv_up2(c)
+
+        x = self.adain1(x, c)
+        x = self.upsample(x)        
+        x = self.dropout(x)
+        x = self.dconv_up1(x)
+        
+        x = self.conv_last(x)
+        out = self.up_last(x)
+        
+        return self.activation(out)
+
+class MLP(nn.Module):
+    def __init__(self, in_f, hidden_dim, out_f):
+        super(MLP, self).__init__()
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.mlp = nn.Sequential(nn.Linear(in_f, hidden_dim),
+                                nn.LeakyReLU(inplace=True),
+                                nn.Linear(hidden_dim, hidden_dim),
+                                nn.LeakyReLU(inplace=True),
+                                nn.Linear(hidden_dim, out_f),)
+
+    def forward(self, x):
+        x = self.pool(x)
+        x = self.mlp(x)
+        return x
+
+class Conv2d1x1(nn.Module):
+    def __init__(self, in_f, hidden_dim, out_f):
+        super(Conv2d1x1, self).__init__()
+        self.conv2d = nn.Sequential(nn.Conv2d(in_f, hidden_dim, 1, 1),
+                                nn.LeakyReLU(inplace=True),
+                                nn.Conv2d(hidden_dim, hidden_dim, 1, 1),
+                                nn.LeakyReLU(inplace=True),
+                                nn.Conv2d(hidden_dim, out_f, 1, 1),)
+
+    def forward(self, x):
+        x = self.conv2d(x)
+        return x
+
+class Head(nn.Module):
+    def __init__(self, in_f, hidden_dim, out_f):
+        super(Head, self).__init__()
+        self.do = nn.Dropout(0.2)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.mlp = nn.Sequential(nn.Linear(in_f, hidden_dim),
+                                nn.LeakyReLU(inplace=True),
+                                nn.Linear(hidden_dim, out_f),)
+
+    def forward(self, x):
+        bs = x.size()[0]
+        x_feat = self.pool(x).view(bs, -1)
+        x = self.mlp(x_feat)
+        x = self.do(x)
+        return x, x_feat
+

loss/__init__.py

@@ -0,0 +1,13 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)
+
+from metrics.registry import LOSSFUNC
+
+from .cross_entropy_loss import CrossEntropyLoss
+from .contrastive_regularization import ContrastiveLoss
+from .l1_loss import L1Loss

loss/abstract_loss_func.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+
+class AbstractLossClass(nn.Module):
+    """Abstract class for loss functions."""
+    def __init__(self):
+        super(AbstractLossClass, self).__init__()
+
+    def forward(self, pred, label):
+        """
+        Args:
+            pred: prediction of the model
+            label: ground truth label
+            
+        Return:
+            loss: loss value
+        """
+        raise NotImplementedError('Each subclass should implement the forward method.')

loss/contrastive_regularization.py

@@ -0,0 +1,78 @@
+import random
+from collections import defaultdict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .abstract_loss_func import AbstractLossClass
+from metrics.registry import LOSSFUNC
+
+
+def swap_spe_features(type_list, value_list):
+    type_list = type_list.cpu().numpy().tolist()
+    # get index
+    index_list = list(range(len(type_list)))
+
+    # init a dict, where its key is the type and value is the index
+    spe_dict = defaultdict(list)
+
+    # do for-loop to get spe dict
+    for i, one_type in enumerate(type_list):
+        spe_dict[one_type].append(index_list[i])
+
+    # shuffle the value list of each key
+    for keys in spe_dict.keys():
+        random.shuffle(spe_dict[keys])
+    
+    # generate a new index list for the value list
+    new_index_list = []
+    for one_type in type_list:
+        value = spe_dict[one_type].pop()
+        new_index_list.append(value)
+
+    # swap the value_list by new_index_list
+    value_list_new = value_list[new_index_list]
+
+    return value_list_new
+
+
+@LOSSFUNC.register_module(module_name="contrastive_regularization")
+class ContrastiveLoss(AbstractLossClass):
+    def __init__(self, margin=1.0):
+        super().__init__()
+        self.margin = margin
+    
+    def contrastive_loss(self, anchor, positive, negative):
+        dist_pos = F.pairwise_distance(anchor, positive)
+        dist_neg = F.pairwise_distance(anchor, negative)
+        # Compute loss as the distance between anchor and negative minus the distance between anchor and positive
+        loss = torch.mean(torch.clamp(dist_pos - dist_neg + self.margin, min=0.0))
+        return loss
+
+    def forward(self, common, specific, spe_label):
+        # prepare
+        bs = common.shape[0]
+        real_common, fake_common = common.chunk(2)
+        ### common real
+        idx_list = list(range(0, bs//2))
+        random.shuffle(idx_list)
+        real_common_anchor = common[idx_list]
+        ### common fake
+        idx_list = list(range(bs//2, bs))
+        random.shuffle(idx_list)
+        fake_common_anchor = common[idx_list]
+        ### specific
+        specific_anchor = swap_spe_features(spe_label, specific)
+        real_specific_anchor, fake_specific_anchor = specific_anchor.chunk(2)
+        real_specific, fake_specific = specific.chunk(2)
+
+        # Compute the contrastive loss of common between real and fake
+        loss_realcommon = self.contrastive_loss(real_common, real_common_anchor, fake_common_anchor)
+        loss_fakecommon = self.contrastive_loss(fake_common, fake_common_anchor, real_common_anchor)
+
+        # Comupte the constrastive loss of specific between real and fake
+        loss_realspecific = self.contrastive_loss(real_specific, real_specific_anchor, fake_specific_anchor)
+        loss_fakespecific = self.contrastive_loss(fake_specific, fake_specific_anchor, real_specific_anchor)
+
+        # Compute the final loss as the sum of all contrastive losses
+        loss = loss_realcommon + loss_fakecommon + loss_fakespecific + loss_realspecific
+        return loss

loss/cross_entropy_loss.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+from .abstract_loss_func import AbstractLossClass
+from metrics.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="cross_entropy")
+class CrossEntropyLoss(AbstractLossClass):
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(self, inputs, targets):
+        """
+        Computes the cross-entropy loss.
+
+        Args:
+            inputs: A PyTorch tensor of size (batch_size, num_classes) containing the predicted scores.
+            targets: A PyTorch tensor of size (batch_size) containing the ground-truth class indices.
+
+        Returns:
+            A scalar tensor representing the cross-entropy loss.
+        """
+        # Compute the cross-entropy loss
+        loss = self.loss_fn(inputs, targets)
+
+        return loss

loss/l1_loss.py

@@ -0,0 +1,19 @@
+import torch.nn as nn
+from .abstract_loss_func import AbstractLossClass
+from metrics.registry import LOSSFUNC
+
+
+@LOSSFUNC.register_module(module_name="l1loss")
+class L1Loss(AbstractLossClass):
+    def __init__(self):
+        super().__init__()
+        self.loss_fn = nn.L1Loss()
+
+    def forward(self, inputs, targets):
+        """
+        Computes the l1 loss.
+        """
+        # Compute the l1 loss
+        loss = self.loss_fn(inputs, targets)
+
+        return loss

metrics/__init__.py

@@ -0,0 +1,7 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)

metrics/base_metrics_class.py

@@ -0,0 +1,205 @@
+import numpy as np
+from sklearn import metrics
+from collections import defaultdict
+import torch
+import torch.nn as nn
+
+
+def get_accracy(output, label):
+    _, prediction = torch.max(output, 1)    # argmax
+    correct = (prediction == label).sum().item()
+    accuracy = correct / prediction.size(0)
+    return accuracy
+
+
+def get_prediction(output, label):
+    prob = nn.functional.softmax(output, dim=1)[:, 1]
+    prob = prob.view(prob.size(0), 1)
+    label = label.view(label.size(0), 1)
+    #print(prob.size(), label.size())
+    datas = torch.cat((prob, label.float()), dim=1)
+    return datas
+
+
+def calculate_metrics_for_train(label, output):
+    if output.size(1) == 2:
+        prob = torch.softmax(output, dim=1)[:, 1]
+    else:
+        prob = output
+
+    # Accuracy
+    _, prediction = torch.max(output, 1)
+    correct = (prediction == label).sum().item()
+    accuracy = correct / prediction.size(0)
+
+    # Average Precision
+    y_true = label.cpu().detach().numpy()
+    y_pred = prob.cpu().detach().numpy()
+    ap = metrics.average_precision_score(y_true, y_pred)
+
+    # AUC and EER
+    try:
+        fpr, tpr, thresholds = metrics.roc_curve(label.squeeze().cpu().numpy(),
+                                                 prob.squeeze().cpu().numpy(),
+                                                 pos_label=1)
+    except:
+        # for the case when we only have one sample
+        return None, None, accuracy, ap
+
+    if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+        # for the case when all the samples within a batch is fake/real
+        auc, eer = None, None
+    else:
+        auc = metrics.auc(fpr, tpr)
+        fnr = 1 - tpr
+        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+
+    return auc, eer, accuracy, ap
+
+
+# ------------ compute average metrics of batches---------------------
+class Metrics_batch():
+    def __init__(self):
+        self.tprs = []
+        self.mean_fpr = np.linspace(0, 1, 100)
+        self.aucs = []
+        self.eers = []
+        self.aps = []
+
+        self.correct = 0
+        self.total = 0
+        self.losses = []
+
+    def update(self, label, output):
+        acc = self._update_acc(label, output)
+        if output.size(1) == 2:
+            prob = torch.softmax(output, dim=1)[:, 1]
+        else:
+            prob = output
+        #label = 1-label
+        #prob = torch.softmax(output, dim=1)[:, 1]
+        auc, eer = self._update_auc(label, prob)
+        ap = self._update_ap(label, prob)
+
+        return acc, auc, eer, ap
+
+    def _update_auc(self, lab, prob):
+        fpr, tpr, thresholds = metrics.roc_curve(lab.squeeze().cpu().numpy(),
+                                                 prob.squeeze().cpu().numpy(),
+                                                 pos_label=1)
+        if np.isnan(fpr[0]) or np.isnan(tpr[0]):
+            return -1, -1
+
+        auc = metrics.auc(fpr, tpr)
+        interp_tpr = np.interp(self.mean_fpr, fpr, tpr)
+        interp_tpr[0] = 0.0
+        self.tprs.append(interp_tpr)
+        self.aucs.append(auc)
+
+        # return auc
+
+        # EER
+        fnr = 1 - tpr
+        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+        self.eers.append(eer)
+
+        return auc, eer
+
+    def _update_acc(self, lab, output):
+        _, prediction = torch.max(output, 1)    # argmax
+        correct = (prediction == lab).sum().item()
+        accuracy = correct / prediction.size(0)
+        # self.accs.append(accuracy)
+        self.correct = self.correct+correct
+        self.total = self.total+lab.size(0)
+        return accuracy
+
+    def _update_ap(self, label, prob):
+        y_true = label.cpu().detach().numpy()
+        y_pred = prob.cpu().detach().numpy()
+        ap = metrics.average_precision_score(y_true,y_pred)
+        self.aps.append(ap)
+
+        return np.mean(ap)
+
+    def get_mean_metrics(self):
+        mean_acc, std_acc = self.correct/self.total, 0
+        mean_auc, std_auc = self._mean_auc()
+        mean_err, std_err = np.mean(self.eers), np.std(self.eers)
+        mean_ap, std_ap = np.mean(self.aps), np.std(self.aps)
+        
+        return {'acc':mean_acc, 'auc':mean_auc, 'eer':mean_err, 'ap':mean_ap}
+
+    def _mean_auc(self):
+        mean_tpr = np.mean(self.tprs, axis=0)
+        mean_tpr[-1] = 1.0
+        mean_auc = metrics.auc(self.mean_fpr, mean_tpr)
+        std_auc = np.std(self.aucs)
+        return mean_auc, std_auc
+
+    def clear(self):
+        self.tprs.clear()
+        self.aucs.clear()
+        # self.accs.clear()
+        self.correct=0
+        self.total=0
+        self.eers.clear()
+        self.aps.clear()
+        self.losses.clear()
+
+
+# ------------ compute average metrics of all data ---------------------
+class Metrics_all():
+    def __init__(self):
+        self.probs = []
+        self.labels = []
+        self.correct = 0
+        self.total = 0
+
+    def store(self, label, output):
+        prob = torch.softmax(output, dim=1)[:, 1]
+        _, prediction = torch.max(output, 1)    # argmax
+        correct = (prediction == label).sum().item()
+        self.correct += correct
+        self.total += label.size(0)
+        self.labels.append(label.squeeze().cpu().numpy())
+        self.probs.append(prob.squeeze().cpu().numpy())
+
+    def get_metrics(self):
+        y_pred = np.concatenate(self.probs)
+        y_true = np.concatenate(self.labels)
+        # auc
+        fpr, tpr, thresholds = metrics.roc_curve(y_true,y_pred,pos_label=1)
+        auc = metrics.auc(fpr, tpr)
+        # eer
+        fnr = 1 - tpr
+        eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+        # ap
+        ap = metrics.average_precision_score(y_true,y_pred)
+        # acc
+        acc = self.correct / self.total
+        return {'acc':acc, 'auc':auc, 'eer':eer, 'ap':ap}
+
+    def clear(self):
+        self.probs.clear()
+        self.labels.clear()
+        self.correct = 0
+        self.total = 0
+
+
+# only used to record a series of scalar value
+class Recorder:
+    def __init__(self):
+        self.sum = 0
+        self.num = 0
+    def update(self, item, num=1):
+        if item is not None:
+            self.sum += item * num
+            self.num += num
+    def average(self):
+        if self.num == 0:
+            return None
+        return self.sum/self.num
+    def clear(self):
+        self.sum = 0
+        self.num = 0

metrics/registry.py

@@ -0,0 +1,20 @@
+class Registry(object):
+    def __init__(self):
+        self.data = {}
+    
+    def register_module(self, module_name=None):
+        def _register(cls):
+            name = module_name
+            if module_name is None:
+                name = cls.__name__
+            self.data[name] = cls
+            return cls
+        return _register
+    
+    def __getitem__(self, key):
+        return self.data[key]
+    
+BACKBONE = Registry()
+DETECTOR = Registry()
+TRAINER  = Registry()
+LOSSFUNC = Registry()

metrics/utils.py

@@ -0,0 +1,88 @@
+from sklearn import metrics
+import numpy as np
+
+
+def parse_metric_for_print(metric_dict):
+    if metric_dict is None:
+        return "\n"
+    str = "\n"
+    str += "================================ Each dataset best metric ================================ \n"
+    for key, value in metric_dict.items():
+        if key != 'avg':
+            str= str+ f"| {key}: "
+            for k,v in value.items():
+                str = str + f" {k}={v} "
+            str= str+ "| \n"
+        else:
+            str += "============================================================================================= \n"
+            str += "================================== Average best metric ====================================== \n"
+            avg_dict = value
+            for avg_key, avg_value in avg_dict.items():
+                if avg_key == 'dataset_dict':
+                    for key,value in avg_value.items():
+                        str = str + f"| {key}: {value} | \n"
+                else:
+                    str = str + f"| avg {avg_key}: {avg_value} | \n"
+    str += "============================================================================================="
+    return str
+
+
+def get_test_metrics(y_pred, y_true, img_names=None, logger=None):
+    def get_video_metrics(image, pred, label):
+        result_dict = {}
+        new_label = []
+        new_pred = []
+        # print(image[0])
+        # print(pred.shape)
+        # print(label.shape)
+        for item in np.transpose(np.stack((image, pred, label)), (1, 0)):
+
+            s = item[0]
+            if '\\' in s:
+                parts = s.split('\\')
+            else:
+                parts = s.split('/')
+            a = parts[-2]
+            b = parts[-1]
+
+            if a not in result_dict:
+                result_dict[a] = []
+
+            result_dict[a].append(item)
+        image_arr = list(result_dict.values())
+
+        for video in image_arr:
+            pred_sum = 0
+            label_sum = 0
+            leng = 0
+            for frame in video:
+                pred_sum += float(frame[1])
+                label_sum += int(frame[2])
+                leng += 1
+            new_pred.append(pred_sum / leng)
+            new_label.append(int(label_sum / leng))
+        fpr, tpr, thresholds = metrics.roc_curve(new_label, new_pred)
+        v_auc = metrics.auc(fpr, tpr)
+        fnr = 1 - tpr
+        v_eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+        return v_auc, v_eer
+
+
+    y_pred = y_pred.squeeze()
+    
+    # For UCF, where labels for different manipulations are not consistent.
+    y_true[y_true >= 1] = 1
+    # auc
+    fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred, pos_label=1)
+    auc = metrics.auc(fpr, tpr)
+    # eer
+    fnr = 1 - tpr
+    eer = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
+    # ap
+    ap = metrics.average_precision_score(y_true, y_pred)
+    # acc
+    prediction_class = (y_pred > 0.5).astype(int)
+    correct = (prediction_class == np.clip(y_true, a_min=0, a_max=1)).sum().item()
+    acc = correct / len(prediction_class)
+
+    return {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap, 'pred': y_pred, 'label': y_true}

networks/__init__.py

@@ -0,0 +1,11 @@
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)
+
+from metrics.registry import BACKBONE
+
+from .xception import Xception

networks/xception.py

@@ -0,0 +1,285 @@
+'''
+# author: Zhiyuan Yan
+# email: zhiyuanyan@link.cuhk.edu.cn
+# date: 2023-0706
+
+The code is mainly modified from GitHub link below:
+https://github.com/ondyari/FaceForensics/blob/master/classification/network/xception.py
+'''
+
+import os
+import argparse
+import logging
+
+import math
+import torch
+# import pretrainedmodels
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch.utils.model_zoo as model_zoo
+from torch.nn import init
+from typing import Union
+from metrics.registry import BACKBONE
+
+logger = logging.getLogger(__name__)
+
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
+        super(SeparableConv2d, self).__init__()
+
+        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size,
+                               stride, padding, dilation, groups=in_channels, bias=bias)
+        self.pointwise = nn.Conv2d(
+            in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pointwise(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, in_filters, out_filters, reps, strides=1, start_with_relu=True, grow_first=True):
+        super(Block, self).__init__()
+
+        if out_filters != in_filters or strides != 1:
+            self.skip = nn.Conv2d(in_filters, out_filters,
+                                  1, stride=strides, bias=False)
+            self.skipbn = nn.BatchNorm2d(out_filters)
+        else:
+            self.skip = None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep = []
+
+        filters = in_filters
+        if grow_first:   # whether the number of filters grows first
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(in_filters, out_filters,
+                                       3, stride=1, padding=1, bias=False))
+            rep.append(nn.BatchNorm2d(out_filters))
+            filters = out_filters
+
+        for i in range(reps-1):
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(filters, filters,
+                                       3, stride=1, padding=1, bias=False))
+            rep.append(nn.BatchNorm2d(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(in_filters, out_filters,
+                                       3, stride=1, padding=1, bias=False))
+            rep.append(nn.BatchNorm2d(out_filters))
+
+        if not start_with_relu:
+            rep = rep[1:]
+        else:
+            rep[0] = nn.ReLU(inplace=False)
+
+        if strides != 1:
+            rep.append(nn.MaxPool2d(3, strides, 1))
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x += skip
+        return x
+
+def add_gaussian_noise(ins, mean=0, stddev=0.2):
+    noise = ins.data.new(ins.size()).normal_(mean, stddev)
+    return ins + noise
+
+
+@BACKBONE.register_module(module_name="xception")
+class Xception(nn.Module):
+    """
+    Xception optimized for the ImageNet dataset, as specified in
+    https://arxiv.org/pdf/1610.02357.pdf
+    """
+
+    def __init__(self, xception_config):
+        """ Constructor
+        Args:
+            xception_config: configuration file with the dict format
+        """
+        super(Xception, self).__init__()
+        self.num_classes = xception_config["num_classes"]
+        self.mode = xception_config["mode"]
+        inc = xception_config["inc"]
+        dropout = xception_config["dropout"]
+
+        # Entry flow
+        self.conv1 = nn.Conv2d(inc, 32, 3, 2, 0, bias=False)
+        
+        self.bn1 = nn.BatchNorm2d(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, bias=False)
+        self.bn2 = nn.BatchNorm2d(64)
+        # do relu here
+
+        self.block1 = Block(
+            64, 128, 2, 2, start_with_relu=False, grow_first=True)
+        self.block2 = Block(
+            128, 256, 2, 2, start_with_relu=True, grow_first=True)
+        self.block3 = Block(
+            256, 728, 2, 2, start_with_relu=True, grow_first=True)
+
+        # middle flow
+        self.block4 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block5 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block6 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block7 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+
+        self.block8 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block9 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block10 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+        self.block11 = Block(
+            728, 728, 3, 1, start_with_relu=True, grow_first=True)
+
+        # Exit flow
+        self.block12 = Block(
+            728, 1024, 2, 2, start_with_relu=True, grow_first=False)
+
+        self.conv3 = SeparableConv2d(1024, 1536, 3, 1, 1)
+        self.bn3 = nn.BatchNorm2d(1536)
+
+        # do relu here
+        self.conv4 = SeparableConv2d(1536, 2048, 3, 1, 1)
+        self.bn4 = nn.BatchNorm2d(2048)
+        # used for iid
+        final_channel = 2048
+        if self.mode == 'adjust_channel_iid':
+            final_channel = 512
+            self.mode = 'adjust_channel'
+        self.last_linear = nn.Linear(final_channel, self.num_classes)
+        if dropout:
+            self.last_linear = nn.Sequential(
+                nn.Dropout(p=dropout),
+                nn.Linear(final_channel, self.num_classes)
+            )
+
+        self.adjust_channel = nn.Sequential(
+            nn.Conv2d(2048, 512, 1, 1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(inplace=False),
+        )
+           
+    def fea_part1_0(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)  
+
+        return x
+
+    def fea_part1_1(self, x):  
+        
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x) 
+
+        return x
+    
+    def fea_part1(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)     
+        
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x) 
+
+        return x
+    
+    def fea_part2(self, x):
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+
+        return x
+
+    def fea_part3(self, x):
+        if self.mode == "shallow_xception":
+            return x
+        else:
+            x = self.block4(x)
+            x = self.block5(x)
+            x = self.block6(x)
+            x = self.block7(x)
+        return x
+
+    def fea_part4(self, x):
+        if self.mode == "shallow_xception":
+            x = self.block12(x)
+        else:
+            x = self.block8(x)
+            x = self.block9(x)
+            x = self.block10(x)
+            x = self.block11(x)
+            x = self.block12(x)
+        return x
+
+    def fea_part5(self, x):
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+
+        return x
+     
+    def features(self, input):
+        x = self.fea_part1(input)    
+
+        x = self.fea_part2(x)
+        x = self.fea_part3(x)
+        x = self.fea_part4(x)
+
+        x = self.fea_part5(x)
+
+        if self.mode == 'adjust_channel':
+            x = self.adjust_channel(x)
+        
+        return x
+
+    def classifier(self, features,id_feat=None):
+        # for iid
+        if self.mode == 'adjust_channel':
+            x = features
+        else:
+            x = self.relu(features)
+
+        if len(x.shape) == 4:
+            x = F.adaptive_avg_pool2d(x, (1, 1))
+            x = x.view(x.size(0), -1)
+        self.last_emb = x
+        # for iid
+        if id_feat!=None:
+            out = self.last_linear(x-id_feat)
+        else:
+            out = self.last_linear(x)
+        return out
+
+    def forward(self, input):
+        x = self.features(input)
+        out = self.classifier(x)
+        return out, x

optimizor/LinearLR.py

@@ -0,0 +1,20 @@
+import torch
+from torch.optim import SGD
+from torch.optim.lr_scheduler import _LRScheduler
+
+class LinearDecayLR(_LRScheduler):
+    def __init__(self, optimizer, n_epoch, start_decay, last_epoch=-1):
+        self.start_decay=start_decay
+        self.n_epoch=n_epoch
+        super(LinearDecayLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        last_epoch = self.last_epoch
+        n_epoch=self.n_epoch
+        b_lr=self.base_lrs[0]
+        start_decay=self.start_decay
+        if last_epoch>start_decay:
+            lr=b_lr-b_lr/(n_epoch-start_decay)*(last_epoch-start_decay)
+        else:
+            lr=b_lr
+        return [lr]

optimizor/SAM.py

@@ -0,0 +1,77 @@
+# borrowed from
+
+import torch
+
+import torch
+import torch.nn as nn
+
+def disable_running_stats(model):
+    def _disable(module):
+        if isinstance(module, nn.BatchNorm2d):
+            module.backup_momentum = module.momentum
+            module.momentum = 0
+
+    model.apply(_disable)
+
+def enable_running_stats(model):
+    def _enable(module):
+        if isinstance(module, nn.BatchNorm2d) and hasattr(module, "backup_momentum"):
+            module.momentum = module.backup_momentum
+
+    model.apply(_enable)
+
+class SAM(torch.optim.Optimizer):
+    def __init__(self, params, base_optimizer, rho=0.05, **kwargs):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = dict(rho=rho, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group["rho"] / (grad_norm + 1e-12)
+
+            for p in group["params"]:
+                if p.grad is None: continue
+                e_w = p.grad * scale.to(p)
+                p.add_(e_w)  # climb to the local maximum "w + e(w)"
+                self.state[p]["e_w"] = e_w
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None: continue
+                p.sub_(self.state[p]["e_w"])  # get back to "w" from "w + e(w)"
+
+        self.base_optimizer.step()  # do the actual "sharpness-aware" update
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
+        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
+
+        self.first_step(zero_grad=True)
+        closure()
+        self.second_step()
+
+    def _grad_norm(self):
+        shared_device = self.param_groups[0]["params"][0].device  # put everything on the same device, in case of model parallelism
+        norm = torch.norm(
+                    torch.stack([
+                        p.grad.norm(p=2).to(shared_device)
+                        for group in self.param_groups for p in group["params"]
+                        if p.grad is not None
+                    ]),
+                    p=2
+               )
+        return norm

trainer/trainer.py

@@ -0,0 +1,441 @@
+# This script was adapted from the DeepfakeBench training code,
+# originally authored by Zhiyuan Yan (zhiyuanyan@link.cuhk.edu.cn)
+
+# Original: https://github.com/SCLBD/DeepfakeBench/blob/main/training/train.py
+
+import os
+import sys
+current_file_path = os.path.abspath(__file__)
+parent_dir = os.path.dirname(os.path.dirname(current_file_path))
+project_root_dir = os.path.dirname(parent_dir)
+sys.path.append(parent_dir)
+sys.path.append(project_root_dir)
+
+import pickle
+import datetime
+import logging
+import numpy as np
+from copy import deepcopy
+from collections import defaultdict
+from tqdm import tqdm
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.nn import DataParallel
+from torch.utils.tensorboard import SummaryWriter
+from metrics.base_metrics_class import Recorder
+from torch.optim.swa_utils import AveragedModel, SWALR
+from torch import distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from sklearn import metrics
+from metrics.utils import get_test_metrics
+
+FFpp_pool=['FaceForensics++','FF-DF','FF-F2F','FF-FS','FF-NT']#
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Trainer(object):
+    def __init__(
+        self,
+        config,
+        model,
+        optimizer,
+        scheduler,
+        logger,
+        metric_scoring='auc',
+        swa_model=None
+        ):
+        # check if all the necessary components are implemented
+        if config is None or model is None or optimizer is None or logger is None:
+            raise ValueError("config, model, optimizier, logger, and tensorboard writer must be implemented")
+
+        self.config = config
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.swa_model = swa_model
+        self.writers = {}  # dict to maintain different tensorboard writers for each dataset and metric
+        self.logger = logger
+        self.metric_scoring = metric_scoring
+        # maintain the best metric of all epochs
+        self.best_metrics_all_time = defaultdict(
+            lambda: defaultdict(lambda: float('-inf')
+            if self.metric_scoring != 'eer' else float('inf'))
+        )
+        self.speed_up()  # move model to GPU
+
+        # create directory path
+        self.log_dir = self.config['log_dir']
+        print("Making dir ", self.log_dir)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+    def get_writer(self, phase, dataset_key, metric_key):
+        phase = phase.split('/')[-1]
+        dataset_key = dataset_key.split('/')[-1]
+        metric_key = metric_key.split('/')[-1]
+        writer_key = f"{phase}-{dataset_key}-{metric_key}"
+        if writer_key not in self.writers:
+            # update directory path
+            writer_path = os.path.join(
+                self.log_dir,
+                phase,
+                dataset_key,
+                metric_key,
+                "metric_board"
+            )
+            os.makedirs(writer_path, exist_ok=True)
+            # update writers dictionary
+            self.writers[writer_key] = SummaryWriter(writer_path)
+        return self.writers[writer_key]
+
+    def speed_up(self):
+        self.model.to(device)
+        self.model.device = device
+        if self.config['ddp'] == True:
+            num_gpus = torch.cuda.device_count()
+            print(f'avai gpus: {num_gpus}')
+            # local_rank=[i for i in range(0,num_gpus)]
+            self.model = DDP(self.model, device_ids=[self.config['local_rank']],find_unused_parameters=True, output_device=self.config['local_rank'])
+            #self.optimizer =  nn.DataParallel(self.optimizer, device_ids=[int(os.environ['LOCAL_RANK'])])
+
+    def setTrain(self):
+        self.model.train()
+        self.train = True
+
+    def setEval(self):
+        self.model.eval()
+        self.train = False
+
+    def load_ckpt(self, model_path):
+        if os.path.isfile(model_path):
+            saved = torch.load(model_path, map_location='cpu')
+            suffix = model_path.split('.')[-1]
+            if suffix == 'p':
+                self.model.load_state_dict(saved.state_dict())
+            else:
+                self.model.load_state_dict(saved)
+            self.logger.info('Model found in {}'.format(model_path))
+        else:
+            raise NotImplementedError(
+                "=> no model found at '{}'".format(model_path))
+
+    def save_ckpt(self, phase, dataset_key,ckpt_info=None):
+        save_dir = self.log_dir
+        os.makedirs(save_dir, exist_ok=True)
+        ckpt_name = f"ckpt_best.pth"
+        save_path = os.path.join(save_dir, ckpt_name)
+        if self.config['ddp'] == True:
+            torch.save(self.model.state_dict(), save_path)
+        else:
+            if 'svdd' in self.config['model_name']:
+                torch.save({'R': self.model.R,
+                            'c': self.model.c,
+                            'state_dict': self.model.state_dict(),}, save_path)
+            else:
+                torch.save(self.model.state_dict(), save_path)
+        self.logger.info(f"Checkpoint saved to {save_path}, current ckpt is {ckpt_info}")
+
+    def save_swa_ckpt(self):
+        save_dir = self.log_dir
+        os.makedirs(save_dir, exist_ok=True)
+        ckpt_name = f"swa.pth"
+        save_path = os.path.join(save_dir, ckpt_name)
+        torch.save(self.swa_model.state_dict(), save_path)
+        self.logger.info(f"SWA Checkpoint saved to {save_path}")
+
+    def save_feat(self, phase, fea, dataset_key):
+        save_dir = os.path.join(self.log_dir, phase, dataset_key)
+        os.makedirs(save_dir, exist_ok=True)
+        features = fea
+        feat_name = f"feat_best.npy"
+        save_path = os.path.join(save_dir, feat_name)
+        np.save(save_path, features)
+        self.logger.info(f"Feature saved to {save_path}")
+
+    def save_data_dict(self, phase, data_dict, dataset_key):
+        save_dir = os.path.join(self.log_dir, phase, dataset_key)
+        os.makedirs(save_dir, exist_ok=True)
+        file_path = os.path.join(save_dir, f'data_dict_{phase}.pickle')
+        with open(file_path, 'wb') as file:
+            pickle.dump(data_dict, file)
+        self.logger.info(f"data_dict saved to {file_path}")
+
+    def save_metrics(self, phase, metric_one_dataset, dataset_key):
+        save_dir = os.path.join(self.log_dir, phase, dataset_key)
+        os.makedirs(save_dir, exist_ok=True)
+        file_path = os.path.join(save_dir, 'metric_dict_best.pickle')
+        with open(file_path, 'wb') as file:
+            pickle.dump(metric_one_dataset, file)
+        self.logger.info(f"Metrics saved to {file_path}")
+
+    def train_step(self,data_dict):
+        if self.config['optimizer']['type']=='sam':
+            for i in range(2):
+                predictions = self.model(data_dict)
+                losses = self.model.get_losses(data_dict, predictions)
+                if i == 0:
+                    pred_first = predictions
+                    losses_first = losses
+                self.optimizer.zero_grad()
+                losses['overall'].backward()
+                if i == 0:
+                    self.optimizer.first_step(zero_grad=True)
+                else:
+                    self.optimizer.second_step(zero_grad=True)
+            return losses_first, pred_first
+        else:
+            predictions = self.model(data_dict)
+            if type(self.model) is DDP:
+                losses = self.model.module.get_losses(data_dict, predictions)
+            else:
+                losses = self.model.get_losses(data_dict, predictions)
+            self.optimizer.zero_grad()
+            losses['overall'].backward()
+            self.optimizer.step()
+
+            return losses,predictions
+
+    def train_epoch(
+        self,
+        epoch,
+        train_data_loader,
+        validation_data_loaders=None
+        ):
+
+        self.logger.info("===> Epoch[{}] start!".format(epoch))
+        if epoch>=1:
+            times_per_epoch = 2
+        else:
+            times_per_epoch = 1
+
+
+        #times_per_epoch=4
+        validation_step = len(train_data_loader) // times_per_epoch    # validate 10 times per epoch
+        step_cnt = epoch * len(train_data_loader)
+        
+        # define training recorder
+        train_recorder_loss = defaultdict(Recorder)
+        train_recorder_metric = defaultdict(Recorder)
+
+        for iteration, data_dict in tqdm(enumerate(train_data_loader),total=len(train_data_loader)):
+            self.setTrain()
+            # more elegant and more scalable way of moving data to GPU
+            for key in data_dict.keys():
+                if data_dict[key]!=None and key!='name':
+                    data_dict[key]=data_dict[key].cuda()
+
+            losses, predictions=self.train_step(data_dict)
+            # update learning rate
+
+            if 'SWA' in self.config and self.config['SWA'] and epoch>self.config['swa_start']:
+                self.swa_model.update_parameters(self.model)
+
+            # compute training metric for each batch data
+            if type(self.model) is DDP:
+                batch_metrics = self.model.module.get_train_metrics(data_dict, predictions)
+            else:
+                batch_metrics = self.model.get_train_metrics(data_dict, predictions)
+
+            # store data by recorder
+            ## store metric
+            for name, value in batch_metrics.items():
+                train_recorder_metric[name].update(value)
+            ## store loss
+            for name, value in losses.items():
+                train_recorder_loss[name].update(value)
+
+            # run tensorboard to visualize the training process
+            if iteration % 300 == 0 and self.config['local_rank']==0:
+                if self.config['SWA'] and (epoch>self.config['swa_start'] or self.config['dry_run']):
+                    self.scheduler.step()
+                # info for loss
+                loss_str = f"Iter: {step_cnt}    "
+                for k, v in train_recorder_loss.items():
+                    v_avg = v.average()
+                    if v_avg == None:
+                        loss_str += f"training-loss, {k}: not calculated"
+                        continue
+                    loss_str += f"training-loss, {k}: {v_avg}    "
+                    # tensorboard-1. loss
+                    processed_train_dataset = [dataset.split('/')[-1] for dataset in self.config['train_dataset']]
+                    processed_train_dataset = ','.join(processed_train_dataset)
+                    writer = self.get_writer('train', processed_train_dataset, k)
+                    writer.add_scalar(f'train_loss/{k}', v_avg, global_step=step_cnt)
+                self.logger.info(loss_str)
+                # info for metric
+                metric_str = f"Iter: {step_cnt}    "
+                for k, v in train_recorder_metric.items():
+                    v_avg = v.average()
+                    if v_avg == None:
+                        metric_str += f"training-metric, {k}: not calculated    "
+                        continue
+                    metric_str += f"training-metric, {k}: {v_avg}    "
+                    # tensorboard-2. metric
+                    processed_train_dataset = [dataset.split('/')[-1] for dataset in self.config['train_dataset']]
+                    processed_train_dataset = ','.join(processed_train_dataset)
+                    writer = self.get_writer('train', processed_train_dataset, k)
+                    writer.add_scalar(f'train_metric/{k}', v_avg, global_step=step_cnt)
+                self.logger.info(metric_str)
+
+                # clear recorder.
+                # Note we only consider the current 300 samples for computing batch-level loss/metric
+                for name, recorder in train_recorder_loss.items():  # clear loss recorder
+                    recorder.clear()
+                for name, recorder in train_recorder_metric.items():  # clear metric recorder
+                    recorder.clear()
+
+            # run validation
+            if (step_cnt+1) % validation_step == 0:
+                if validation_data_loaders is not None and ((not self.config['ddp']) or (self.config['ddp'] and dist.get_rank() == 0)):
+                    self.logger.info("===> Validation start!")
+                    validation_best_metric = self.eval(
+                        eval_data_loaders=validation_data_loaders,
+                        eval_stage="validation",
+                        step=step_cnt,
+                        epoch=epoch,
+                        iteration=iteration
+                    )
+                else:
+                    validation_best_metric = None
+
+            step_cnt += 1
+
+            for key in data_dict.keys():
+                if data_dict[key]!=None and key!='name':
+                    data_dict[key]=data_dict[key].cpu()
+        return validation_best_metric
+    
+    def get_respect_acc(self,prob,label):
+        pred = np.where(prob > 0.5, 1, 0)
+        judge = (pred == label)
+        zero_num = len(label) - np.count_nonzero(label)
+        acc_fake = np.count_nonzero(judge[zero_num:]) / len(judge[zero_num:])
+        acc_real = np.count_nonzero(judge[:zero_num]) / len(judge[:zero_num])
+        return acc_real,acc_fake
+
+    def eval_one_dataset(self, data_loader):
+        # define eval recorder
+        eval_recorder_loss = defaultdict(Recorder)
+        prediction_lists = []
+        feature_lists=[]
+        label_lists = []
+        for i, data_dict in tqdm(enumerate(data_loader),total=len(data_loader)):
+            # get data
+            if 'label_spe' in data_dict:
+                data_dict.pop('label_spe')  # remove the specific label
+            data_dict['label'] = torch.where(data_dict['label']!=0, 1, 0)  # fix the label to 0 and 1 only
+            # move data to GPU elegantly
+            for key in data_dict.keys():
+                if data_dict[key]!=None:
+                    data_dict[key]=data_dict[key].cuda()
+            # model forward without considering gradient computation
+            predictions = self.inference(data_dict) #dict with keys cls, feat
+            
+            label_lists += list(data_dict['label'].cpu().detach().numpy())
+            # Get the predicted class for each sample in the batch
+            _, predicted_classes = torch.max(predictions['cls'], dim=1)
+            # Convert the predicted class indices to a list and add to prediction_lists
+            prediction_lists += predicted_classes.cpu().detach().numpy().tolist()
+            feature_lists += list(predictions['feat'].cpu().detach().numpy())
+            if type(self.model) is not AveragedModel:
+                # compute all losses for each batch data
+                if type(self.model) is DDP:
+                    losses = self.model.module.get_losses(data_dict, predictions)
+                else:
+                    losses = self.model.get_losses(data_dict, predictions)
+
+                # store data by recorder
+                for name, value in losses.items():
+                    eval_recorder_loss[name].update(value)
+        return eval_recorder_loss, np.array(prediction_lists), np.array(label_lists),np.array(feature_lists)
+
+    def save_best(self,epoch,iteration,step,losses_one_dataset_recorder,key,metric_one_dataset,eval_stage):
+        best_metric = self.best_metrics_all_time[key].get(self.metric_scoring,
+                                                          float('-inf') if self.metric_scoring != 'eer' else float(
+                                                              'inf'))
+        # Check if the current score is an improvement
+        improved = (metric_one_dataset[self.metric_scoring] > best_metric) if self.metric_scoring != 'eer' else (
+                    metric_one_dataset[self.metric_scoring] < best_metric)
+        if improved:
+            # Update the best metric
+            self.best_metrics_all_time[key][self.metric_scoring] = metric_one_dataset[self.metric_scoring]
+            if key == 'avg':
+                self.best_metrics_all_time[key]['dataset_dict'] = metric_one_dataset['dataset_dict']
+            # Save checkpoint, feature, and metrics if specified in config
+            if eval_stage=='validation' and self.config['save_ckpt'] and key not in FFpp_pool:
+                self.save_ckpt(eval_stage, key, f"{epoch}+{iteration}")
+            self.save_metrics(eval_stage, metric_one_dataset, key)
+        if losses_one_dataset_recorder is not None:
+            # info for each dataset
+            loss_str = f"dataset: {key}    step: {step}    "
+            for k, v in losses_one_dataset_recorder.items():
+                writer = self.get_writer(eval_stage, key, k)
+                v_avg = v.average()
+                if v_avg == None:
+                    print(f'{k} is not calculated')
+                    continue
+                # tensorboard-1. loss
+                writer.add_scalar(f'{eval_stage}_losses/{k}', v_avg, global_step=step)
+                loss_str += f"{eval_stage}-loss, {k}: {v_avg}    "
+            self.logger.info(loss_str)
+        # tqdm.write(loss_str)
+        metric_str = f"dataset: {key}    step: {step}    "
+        for k, v in metric_one_dataset.items():
+            if k == 'pred' or k == 'label' or k=='dataset_dict':
+                continue
+            metric_str += f"{eval_stage}-metric, {k}: {v}    "
+            # tensorboard-2. metric
+            writer = self.get_writer(eval_stage, key, k)
+            writer.add_scalar(f'{eval_stage}_metrics/{k}', v, global_step=step)
+        if 'pred' in metric_one_dataset:
+            acc_real, acc_fake = self.get_respect_acc(metric_one_dataset['pred'], metric_one_dataset['label'])
+            metric_str += f'{eval_stage}-metric, acc_real:{acc_real}; acc_fake:{acc_fake}'
+            writer.add_scalar(f'{eval_stage}_metrics/acc_real', acc_real, global_step=step)
+            writer.add_scalar(f'{eval_stage}_metrics/acc_fake', acc_fake, global_step=step)
+        self.logger.info(metric_str)
+
+    def eval(self, eval_data_loaders, eval_stage, step=None, epoch=None, iteration=None):
+        # set model to eval mode
+        self.setEval()
+
+        # define eval recorder
+        losses_all_datasets = {}
+        metrics_all_datasets = {}
+        best_metrics_per_dataset = defaultdict(dict)  # best metric for each dataset, for each metric
+        avg_metric = {'acc': 0, 'auc': 0, 'eer': 0, 'ap': 0,'dataset_dict':{}} #'video_auc': 0
+        keys = eval_data_loaders.keys()
+        for key in keys:
+            # compute loss for each dataset
+            losses_one_dataset_recorder, predictions_nps, label_nps, feature_nps = self.eval_one_dataset(eval_data_loaders[key])
+            losses_all_datasets[key] = losses_one_dataset_recorder
+            metric_one_dataset=get_test_metrics(y_pred=predictions_nps,y_true=label_nps, logger=self.logger)
+            
+            for metric_name, value in metric_one_dataset.items():
+                if metric_name in avg_metric:
+                    avg_metric[metric_name]+=value
+            avg_metric['dataset_dict'][key] = metric_one_dataset[self.metric_scoring]
+            if type(self.model) is AveragedModel:
+                metric_str = f"Iter Final for SWA:    "
+                for k, v in metric_one_dataset.items():
+                    metric_str += f"{eval_stage}-metric, {k}: {v}    "
+                self.logger.info(metric_str)
+                continue
+            self.save_best(epoch,iteration,step,losses_one_dataset_recorder,key,metric_one_dataset,eval_stage)
+
+        if len(keys)>0 and self.config.get('save_avg',False):
+            # calculate avg value
+            for key in avg_metric:
+                if key != 'dataset_dict':
+                    avg_metric[key] /= len(keys)
+            self.save_best(epoch, iteration, step, None, 'avg', avg_metric, eval_stage)
+
+        self.logger.info(f'===> {eval_stage} Done!')
+        return self.best_metrics_all_time  # return all types of mean metrics for determining the best ckpt
+
+    
+    @torch.no_grad()
+    def inference(self, data_dict):
+        predictions = self.model(data_dict, inference=True)
+        return predictions