train_detector.py

# 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()