Delete test_deps

test_deps/README.md

@@ -1,14 +0,0 @@
-## UCF
-
-This model has been adapted from [DeepfakeBench](https://github.com/SCLBD/DeepfakeBench).
-
-## 
-
-- **Train UCF model**:
-   - Use `train_ucf.py`, which will download necessary pretrained `xception` backbone weights from HuggingFace (if not present locally) and start a training job with logging outputs in `.logs/`.
-   - Customize the training job by editing `config/ucf.yaml`
-     - `pm2 start train_ucf.py --no-autorestart` to train a generalist detector on datasets from `DATASET_META`
-     - `pm2 start train_ucf.py --no-autorestart -- --faces_only` to train a face expert detector on preprocessed-face only datasets
-
-- **Miner Neurons**:
-   - The `UCF` class in `pretrained_ucf.py` is used by miner neurons to load and perform inference with pretrained UCF model weights.

test_deps/config/__init__.py

@@ -1,7 +0,0 @@
-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)

test_deps/config/constants.py

@@ -1,15 +0,0 @@
-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"))

test_deps/config/pretrained_config.yaml

@@ -1,94 +0,0 @@
-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

test_deps/config/pretrained_face_config.yaml

@@ -1,94 +0,0 @@
-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

test_deps/config/train_config.yaml

@@ -1,9 +0,0 @@
-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/

test_deps/config/ucf.yaml

@@ -1,73 +0,0 @@
-# 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

test_deps/config/xception.yaml

@@ -1,86 +0,0 @@
-# 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

test_deps/detectors/__init__.py

@@ -1,11 +0,0 @@
-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

test_deps/detectors/base_detector.py

@@ -1,71 +0,0 @@
-# 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

test_deps/detectors/ucf_detector.py

@@ -1,472 +0,0 @@
-'''
-# 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
-

test_deps/logger.py

@@ -1,36 +0,0 @@
-import os
-import logging
-
-import torch.distributed as dist
-
-class RankFilter(logging.Filter):
-    def __init__(self, rank):
-        super().__init__()
-        self.rank = rank
-
-    def filter(self, record):
-        return dist.get_rank() == self.rank
-
-def create_logger(log_path):
-    # Create log path
-    if os.path.isdir(os.path.dirname(log_path)):
-        os.makedirs(os.path.dirname(log_path), exist_ok=True)
-
-    # Create logger object
-    logger = logging.getLogger()
-    logger.setLevel(logging.INFO)
-    # Create file handler and set the formatter
-    fh = logging.FileHandler(log_path)
-    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-    fh.setFormatter(formatter)
-
-    # Add the file handler to the logger
-    logger.addHandler(fh)
-
-    # Add a stream handler to print to console
-    sh = logging.StreamHandler()
-    sh.setLevel(logging.INFO)  # Set logging level for stream handler
-    sh.setFormatter(formatter)
-    logger.addHandler(sh)
-
-    return logger

test_deps/loss/__init__.py

@@ -1,13 +0,0 @@
-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

test_deps/loss/abstract_loss_func.py

@@ -1,17 +0,0 @@
-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.')

test_deps/loss/contrastive_regularization.py

@@ -1,78 +0,0 @@
-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

test_deps/loss/cross_entropy_loss.py

@@ -1,26 +0,0 @@
-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

test_deps/loss/l1_loss.py

@@ -1,19 +0,0 @@
-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

test_deps/metrics/__init__.py

@@ -1,7 +0,0 @@
-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)

test_deps/metrics/base_metrics_class.py

@@ -1,205 +0,0 @@
-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

test_deps/metrics/registry.py

@@ -1,20 +0,0 @@
-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()

test_deps/metrics/utils.py

@@ -1,88 +0,0 @@
-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}

test_deps/networks/__init__.py

@@ -1,11 +0,0 @@
-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

test_deps/networks/xception.py

@@ -1,285 +0,0 @@
-'''
-# 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

test_deps/optimizor/LinearLR.py

@@ -1,20 +0,0 @@
-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]

test_deps/optimizor/SAM.py

@@ -1,77 +0,0 @@
-# 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

test_deps/train_detector.py

@@ -1,460 +0,0 @@
-# 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
-
-# BitMind's modifications include adding a testing phase, changing the 
-# data load/split pipeline to work with subnet 34's image augmentations
-# and datasets from BitMind HuggingFace repositories, quality of life CLI args,
-# logging changes, etc.
-
-import os
-import sys
-import argparse
-from os.path import join
-import random
-import datetime
-import time
-import yaml
-from tqdm import tqdm
-import numpy as np
-from datetime import timedelta
-from copy import deepcopy
-from PIL import Image as pil_image
-from pathlib import Path
-import gc
-
-import torch
-import torch.nn as nn
-import torch.nn.parallel
-import torch.backends.cudnn as cudnn
-import torch.utils.data
-import torch.optim as optim
-from torch.utils.data.distributed import DistributedSampler
-import torch.distributed as dist
-from torch.utils.data import DataLoader
-
-from optimizor.SAM import SAM
-from optimizor.LinearLR import LinearDecayLR
-
-from trainer.trainer import Trainer
-from arena.detectors.UCF.detectors import DETECTOR
-from metrics.utils import parse_metric_for_print
-from logger import create_logger, RankFilter
-
-from huggingface_hub import hf_hub_download
-
-# BitMind imports (not from original Deepfake Bench repo)
-from bitmind.dataset_processing.load_split_data import load_datasets, create_real_fake_datasets
-from bitmind.image_transforms import base_transforms, random_aug_transforms
-from bitmind.constants import DATASET_META, FACE_TRAINING_DATASET_META
-from config.constants import (
-    CONFIG_PATH,
-    WEIGHTS_DIR,
-    HF_REPO,
-    BACKBONE_CKPT
-)
-
-parser = argparse.ArgumentParser(description='Process some paths.')
-parser.add_argument('--detector_path', type=str, default=CONFIG_PATH, help='path to detector YAML file')
-parser.add_argument('--faces_only', dest='faces_only', action='store_true', default=False)
-parser.add_argument('--no-save_ckpt', dest='save_ckpt', action='store_false', default=True)
-parser.add_argument('--no-save_feat', dest='save_feat', action='store_false', default=True)
-parser.add_argument("--ddp", action='store_true', default=False)
-parser.add_argument('--local_rank', type=int, default=0)
-parser.add_argument('--workers', type=int, default=os.cpu_count() - 1,
-                    help='number of workers for data loading')
-parser.add_argument('--epochs', type=int, default=None, help='number of training epochs')
-
-args = parser.parse_args()
-torch.cuda.set_device(args.local_rank)
-print(f"torch.cuda.device(0): {torch.cuda.device(0)}")
-print(f"torch.cuda.get_device_name(0): {torch.cuda.get_device_name(0)}")
-
-def ensure_backbone_is_available(logger,
-                                 weights_dir=WEIGHTS_DIR,
-                                 model_filename=BACKBONE_CKPT,
-                                 hugging_face_repo_name=HF_REPO):
-    
-    destination_path = Path(weights_dir) / Path(model_filename)
-    if not destination_path.parent.exists():
-        destination_path.parent.mkdir(parents=True, exist_ok=True)
-        logger.info(f"Created directory {destination_path.parent}.")
-    if not destination_path.exists():
-        model_path = hf_hub_download(hugging_face_repo_name, model_filename)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = torch.load(model_path, map_location=device)
-        torch.save(model, destination_path)
-        del model
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        logger.info(f"Downloaded backbone {model_filename} to {destination_path}.")
-    else:
-        logger.info(f"{model_filename} backbone already present at {destination_path}.")
-
-def init_seed(config):
-    if config['manualSeed'] is None:
-        config['manualSeed'] = random.randint(1, 10000)
-    random.seed(config['manualSeed'])
-    if config['cuda']:
-        torch.manual_seed(config['manualSeed'])
-        torch.cuda.manual_seed_all(config['manualSeed'])
-
-def custom_collate_fn(batch):
-    images, labels, source_labels = zip(*batch)
-    
-    images = torch.stack(images, dim=0)  # Stack image tensors into a single tensor
-    labels = torch.LongTensor(labels) 
-    source_labels = torch.LongTensor(source_labels) 
-    
-    data_dict = {
-        'image': images,
-        'label': labels,
-        'label_spe': source_labels,
-        'landmark': None,
-        'mask': None
-    }    
-    return data_dict
-
-def prepare_datasets(config, logger):
-    start_time = log_start_time(logger, "Loading and splitting individual datasets")
-    
-    real_datasets, fake_datasets = load_datasets(dataset_meta=config['dataset_meta'], 
-                                                 expert=config['faces_only'],
-                                                 split_transforms=config['split_transforms'])
-
-    log_finish_time(logger, "Loading and splitting individual datasets", start_time)
-    
-    start_time = log_start_time(logger, "Creating real fake dataset splits")
-    train_dataset, val_dataset, test_dataset = \
-    create_real_fake_datasets(real_datasets,
-                              fake_datasets,
-                              config['split_transforms']['train']['transform'],
-                              config['split_transforms']['validation']['transform'],
-                              config['split_transforms']['test']['transform'],
-                              source_labels=True)
-
-    log_finish_time(logger, "Creating real fake dataset splits", start_time)
-    
-    train_loader = torch.utils.data.DataLoader(train_dataset,
-                                               batch_size=config['train_batchSize'],
-                                               shuffle=True,
-                                               num_workers=config['workers'],
-                                               drop_last=True,
-                                               collate_fn=custom_collate_fn)
-    val_loader = torch.utils.data.DataLoader(val_dataset,
-                                             batch_size=config['train_batchSize'],
-                                             shuffle=True,
-                                             num_workers=config['workers'],
-                                             drop_last=True,
-                                             collate_fn=custom_collate_fn)
-    test_loader = torch.utils.data.DataLoader(test_dataset,
-                                              batch_size=config['train_batchSize'],
-                                              shuffle=True, 
-                                              num_workers=config['workers'],
-                                              drop_last=True,
-                                              collate_fn=custom_collate_fn)
-
-    print(f"Train size: {len(train_loader.dataset)}")
-    print(f"Validation size: {len(val_loader.dataset)}")
-    print(f"Test size: {len(test_loader.dataset)}")
-
-    return train_loader, val_loader, test_loader
-
-def choose_optimizer(model, config):
-    opt_name = config['optimizer']['type']
-    if opt_name == 'sgd':
-        optimizer = optim.SGD(
-            params=model.parameters(),
-            lr=config['optimizer'][opt_name]['lr'],
-            momentum=config['optimizer'][opt_name]['momentum'],
-            weight_decay=config['optimizer'][opt_name]['weight_decay']
-        )
-        return optimizer
-    elif opt_name == 'adam':
-        optimizer = optim.Adam(
-            params=model.parameters(),
-            lr=config['optimizer'][opt_name]['lr'],
-            weight_decay=config['optimizer'][opt_name]['weight_decay'],
-            betas=(config['optimizer'][opt_name]['beta1'], config['optimizer'][opt_name]['beta2']),
-            eps=config['optimizer'][opt_name]['eps'],
-            amsgrad=config['optimizer'][opt_name]['amsgrad'],
-        )
-        return optimizer
-    elif opt_name == 'sam':
-        optimizer = SAM(
-            model.parameters(), 
-            optim.SGD, 
-            lr=config['optimizer'][opt_name]['lr'],
-            momentum=config['optimizer'][opt_name]['momentum'],
-        )
-    else:
-        raise NotImplementedError('Optimizer {} is not implemented'.format(config['optimizer']))
-    return optimizer
-
-
-def choose_scheduler(config, optimizer):
-    if config['lr_scheduler'] is None:
-        return None
-    elif config['lr_scheduler'] == 'step':
-        scheduler = optim.lr_scheduler.StepLR(
-            optimizer,
-            step_size=config['lr_step'],
-            gamma=config['lr_gamma'],
-        )
-        return scheduler
-    elif config['lr_scheduler'] == 'cosine':
-        scheduler = optim.lr_scheduler.CosineAnnealingLR(
-            optimizer,
-            T_max=config['lr_T_max'],
-            eta_min=config['lr_eta_min'],
-        )
-        return scheduler
-    elif config['lr_scheduler'] == 'linear':
-        scheduler = LinearDecayLR(
-            optimizer,
-            config['nEpochs'],
-            int(config['nEpochs']/4),
-        )
-    else:
-        raise NotImplementedError('Scheduler {} is not implemented'.format(config['lr_scheduler']))
-
-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 log_start_time(logger, process_name):
-    """Log the start time of a process."""
-    start_time = time.time()
-    logger.info(f"{process_name} Start Time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}")
-    return start_time
-
-def log_finish_time(logger, process_name, start_time):
-    """Log the finish time and elapsed time of a process."""
-    finish_time = time.time()
-    elapsed_time = finish_time - start_time
-
-    # Convert elapsed time into hours, minutes, and seconds
-    hours, rem = divmod(elapsed_time, 3600)
-    minutes, seconds = divmod(rem, 60)
-
-    # Log the finish time and elapsed time
-    logger.info(f"{process_name} Finish Time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(finish_time))}")
-    logger.info(f"{process_name} Elapsed Time: {int(hours)} hours, {int(minutes)} minutes, {seconds:.2f} seconds")
-
-def save_config(config, outputs_dir):
-    """
-    Saves a config dictionary as both a pickle file and a YAML file, ensuring only basic types are saved.
-    Also, lists like 'mean' and 'std' are saved in flow style (on a single line).
-    
-    Args:
-        config (dict): The configuration dictionary to save.
-        outputs_dir (str): The directory path where the files will be saved.
-    """
-
-    def is_basic_type(value):
-        """
-        Check if a value is a basic data type that can be saved in YAML.
-        Basic types include int, float, str, bool, list, and dict.
-        """
-        return isinstance(value, (int, float, str, bool, list, dict, type(None)))
-
-    def filter_dict(data_dict):
-        """
-        Recursively filter out any keys from the dictionary whose values contain non-basic types (e.g., objects).
-        """
-        if not isinstance(data_dict, dict):
-            return data_dict
-        
-        filtered_dict = {}
-        for key, value in data_dict.items():
-            if isinstance(value, dict):
-                # Recursively filter nested dictionaries
-                nested_dict = filter_dict(value)
-                if nested_dict:  # Only add non-empty dictionaries
-                    filtered_dict[key] = nested_dict
-            elif is_basic_type(value):
-                # Add if the value is a basic type
-                filtered_dict[key] = value
-            else:
-                # Skip the key if the value is not a basic type (e.g., an object)
-                print(f"Skipping key '{key}' because its value is of type {type(value)}")
-        
-        return filtered_dict
-
-    def save_dict_to_yaml(data_dict, file_path):
-        """
-        Saves a dictionary to a YAML file, excluding any keys where the value is an object or contains an object.
-        Additionally, ensures that specific lists (like 'mean' and 'std') are saved in flow style.
-        
-        Args:
-            data_dict (dict): The dictionary to save.
-            file_path (str): The local file path where the YAML file will be saved.
-        """
-        
-        # Custom representer for lists to force flow style (compact lists)
-        class FlowStyleList(list):
-            pass
-        
-        def flow_style_list_representer(dumper, data):
-            return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)
-        
-        yaml.add_representer(FlowStyleList, flow_style_list_representer)
-
-        # Preprocess specific lists to be in flow style
-        if 'mean' in data_dict:
-            data_dict['mean'] = FlowStyleList(data_dict['mean'])
-        if 'std' in data_dict:
-            data_dict['std'] = FlowStyleList(data_dict['std'])
-
-        try:
-            # Filter the dictionary
-            filtered_dict = filter_dict(data_dict)
-            
-            # Save the filtered dictionary as YAML
-            with open(file_path, 'w') as f:
-                yaml.dump(filtered_dict, f, default_flow_style=False)  # Save with default block style except for FlowStyleList
-            print(f"Filtered dictionary successfully saved to {file_path}")
-        except Exception as e:
-            print(f"Error saving dictionary to YAML: {e}")
-
-    # Save as YAML
-    save_dict_to_yaml(config, outputs_dir + '/config.yaml')
-
-def main():
-    torch.cuda.empty_cache()
-    gc.collect()
-    # parse options and load config
-    with open(args.detector_path, 'r') as f:
-        config = yaml.safe_load(f)
-    with open(os.getcwd() + '/config/train_config.yaml', 'r') as f:
-        config2 = yaml.safe_load(f)
-    if 'label_dict' in config:
-        config2['label_dict']=config['label_dict']
-    config.update(config2)
-
-    config['workers'] = args.workers
-    
-    config['local_rank']=args.local_rank
-    if config['dry_run']:
-        config['nEpochs'] = 0
-        config['save_feat']=False
-
-    if args.epochs: config['nEpochs'] = args.epochs
-
-    config['split_transforms'] = {'train': {'name': 'base_transforms',
-                                            'transform': base_transforms},
-                                  'validation': {'name': 'base_transforms',
-                                                 'transform': base_transforms},
-                                  'test': {'name': 'base_transforms',
-                                           'transform': base_transforms}}
-    config['faces_only'] = args.faces_only
-    config['dataset_meta'] = FACE_TRAINING_DATASET_META if config['faces_only'] else DATASET_META
-    dataset_names = [item["path"] for datasets in config['dataset_meta'].values() for item in datasets]
-    config['train_dataset'] = dataset_names
-    config['save_ckpt'] = args.save_ckpt
-    config['save_feat'] = args.save_feat
-
-    config['specific_task_number'] = len(config['dataset_meta']["fake"]) + 1
-    
-    if config['lmdb']:
-        config['dataset_json_folder'] = 'preprocessing/dataset_json_v3'
-        
-    # create logger
-    timenow=datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
-    
-    outputs_dir =  os.path.join(
-                config['log_dir'],
-                config['model_name'] + '_' + timenow
-            )
-    
-    os.makedirs(outputs_dir, exist_ok=True)
-    logger = create_logger(os.path.join(outputs_dir, 'training.log'))
-    config['log_dir'] = outputs_dir
-    logger.info('Save log to {}'.format(outputs_dir))
-    
-    config['ddp']= args.ddp
-
-    # init seed
-    init_seed(config)
-
-    # set cudnn benchmark if needed
-    if config['cudnn']:
-        cudnn.benchmark = True
-    if config['ddp']:
-        # dist.init_process_group(backend='gloo')
-        dist.init_process_group(
-            backend='nccl',
-            timeout=timedelta(minutes=30)
-        )
-        logger.addFilter(RankFilter(0))
-
-    ensure_backbone_is_available(logger=logger,
-                                 model_filename=config['pretrained'].split('/')[-1],
-                                 hugging_face_repo_name='bitmind/' + config['model_name'])
-    
-    # prepare the model (detector)
-    model_class = DETECTOR[config['model_name']]
-    model = model_class(config)
-    
-    # prepare the optimizer
-    optimizer = choose_optimizer(model, config)
-
-    # prepare the scheduler
-    scheduler = choose_scheduler(config, optimizer)
-
-    # prepare the metric
-    metric_scoring = choose_metric(config)
-
-    # prepare the trainer
-    trainer = Trainer(config, model, optimizer, scheduler, logger, metric_scoring)
-
-    # prepare the data loaders
-    train_loader, val_loader, test_loader = prepare_datasets(config, logger)
-
-    # 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)
-    
-    # save training configs
-    save_config(config, outputs_dir)
-    
-    # start training
-    start_time = log_start_time(logger, "Training")
-    for epoch in range(config['start_epoch'], config['nEpochs'] + 1):
-        trainer.model.epoch = epoch
-        best_metric = trainer.train_epoch(
-                    epoch,
-                    train_data_loader=train_loader,
-                    validation_data_loaders={'val':val_loader}
-                )
-        if best_metric is not None:
-            logger.info(f"===> Epoch[{epoch}] end with validation {metric_scoring}: {parse_metric_for_print(best_metric)}!")
-    logger.info("Stop Training on best Validation metric {}".format(parse_metric_for_print(best_metric))) 
-    log_finish_time(logger, "Training", start_time)
-   
-    # test
-    start_time = log_start_time(logger, "Test")
-    trainer.eval(eval_data_loaders={'test':test_loader}, eval_stage="test")
-    log_finish_time(logger, "Test", start_time)
-    
-    # update
-    if 'svdd' in config['model_name']:
-        model.update_R(epoch)
-    if scheduler is not None:
-        scheduler.step()
-
-    # close the tensorboard writers
-    for writer in trainer.writers.values():
-        writer.close()
-
-    torch.cuda.empty_cache()
-    gc.collect()
-
-if __name__ == '__main__':
-    main()

test_deps/trainer/trainer.py

@@ -1,441 +0,0 @@
-# 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