training/detectors/effort_detector.py

import os
import math
import datetime
import logging
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, calculate_acc_for_train

from .base_detector import AbstractDetector
from detectors import DETECTOR
from networks import BACKBONE
from loss import LOSSFUNC

import loralib as lora
from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig

logger = logging.getLogger(__name__)


@DETECTOR.register_module(module_name='effort')
class EffortDetector(nn.Module):
    def __init__(self, config=None):
        super(EffortDetector, self).__init__()
        self.config = config
        self.backbone = self.build_backbone(config)
        self.head = nn.Linear(1024, config['backbone_config']['num_classes'])
        self.loss_func = nn.CrossEntropyLoss()
        self.prob, self.label = [], []
        self.correct, self.total = 0, 0

    def build_backbone(self, config):
        # Download model
        # https://huggingface.co/openai/clip-vit-large-patch14

        # mean: [0.48145466, 0.4578275, 0.40821073]
        # std: [0.26862954, 0.26130258, 0.27577711]

        # ViT-L/14 224*224
        clip_model = CLIPModel.from_pretrained(self.config["pretrained"])

        # Apply SVD to self_attn layers only
        # ViT-L/14 224*224: 1024-1
        clip_model.vision_model = apply_svd_residual_to_self_attn(clip_model.vision_model, r=1024-64)

        for name, param in clip_model.vision_model.named_parameters():
            print('{}: {}'.format(name, param.requires_grad))
        num_param = sum(p.numel() for p in clip_model.vision_model.parameters() if p.requires_grad)
        num_total_param = sum(p.numel() for p in clip_model.vision_model.parameters())
        print('Number of total parameters: {}, tunable parameters: {}'.format(num_total_param, num_param))

        return clip_model.vision_model

    def features(self, data_dict: dict) -> torch.tensor:
        # data_dict['image']: torch.Size([32, 3, 224, 224])
        feat = self.backbone(data_dict['image'])['pooler_output']
        # feat torch.Size([32, 1024])
        return feat

    def classifier(self, features: torch.tensor) -> torch.tensor:
        return self.head(features)

    # def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
    #     label = data_dict['label']
    #     pred = pred_dict['cls']
    #     loss = self.loss_func(pred, label)

    #     # Regularization term
    #     lambda_reg = 0.1
    #     orthogonal_losses = []
    #     for module in self.backbone.modules():
    #         if isinstance(module, SVDResidualLinear):
    #             # Apply orthogonal constraints to the U_residual and V_residual matrix
    #             orthogonal_losses.append(module.compute_orthogonal_loss())

    #     if orthogonal_losses:
    #         reg_term = sum(orthogonal_losses)
    #         loss += lambda_reg * reg_term

    #     loss_dict = {'overall': loss}
    #     return loss_dict

    def compute_weight_loss(self):
        weight_sum_dict = {}
        num_weight_dict = {}
        for name, module in self.backbone.named_modules():
            if isinstance(module, SVDResidualLinear):
                weight_curr = module.compute_current_weight()
                if str(weight_curr.size()) not in weight_sum_dict.keys():
                    weight_sum_dict[str(weight_curr.size())] = weight_curr
                    num_weight_dict[str(weight_curr.size())] = 1
                else:
                    weight_sum_dict[str(weight_curr.size())] += weight_curr
                    num_weight_dict[str(weight_curr.size())] += 1

        loss2 = 0.0
        for k in weight_sum_dict.keys():
            _, S_sum, _ = torch.linalg.svd(weight_sum_dict[k], full_matrices=False)
            loss2 += -torch.mean(S_sum)
        loss2 /= len(weight_sum_dict.keys())
        return loss2

    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
        label = data_dict['label']  # Tensor of shape [batch_size]
        pred = pred_dict['cls']     # Tensor of shape [batch_size, num_classes]

        # Compute overall loss using all samples
        loss = self.loss_func(pred, label)

        # Create masks for real and fake classes
        mask_real = label == 0  # Boolean tensor
        mask_fake = label == 1  # Boolean tensor

        # Compute loss for real class
        if mask_real.sum() > 0:
            pred_real = pred[mask_real]
            label_real = label[mask_real]
            loss_real = self.loss_func(pred_real, label_real)
        else:
            # No real samples in batch
            loss_real = torch.tensor(0.0, device=pred.device)

        # Compute loss for fake class
        if mask_fake.sum() > 0:
            pred_fake = pred[mask_fake]
            label_fake = label[mask_fake]
            loss_fake = self.loss_func(pred_fake, label_fake)
        else:
            # No fake samples in batch
            loss_fake = torch.tensor(0.0, device=pred.device)

        # loss2 = self.compute_weight_loss()
        # overall_loss = loss + loss2

        # Return a dictionary with all losses
        loss_dict = {
            'overall': loss,
            'real_loss': loss_real,
            'fake_loss': loss_fake,
            # 'erank_loss': loss2
        }
        return loss_dict

    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
        label = data_dict['label']
        pred = pred_dict['cls']

        # compute metrics for batch data
        # auc, eer, acc, ap = calculate_metrics_for_train(label.detach(), pred.detach())
        # metric_batch_dict = {'acc': acc, 'auc': auc, 'eer': eer, 'ap': ap}

        acc, mAP = calculate_acc_for_train(label.detach(), pred.detach(), self.config['backbone_config']['num_classes'])
        metric_batch_dict = {'acc': acc, 'mAP': mAP}

        return metric_batch_dict


    def forward(self, data_dict: dict, inference=False) -> dict:
        # get the features by backbone
        features = self.features(data_dict)
        # get the prediction by classifier
        pred = self.classifier(features)
        # get the probability of the pred
        # prob = torch.softmax(pred, dim=1)[:, 1]
        prob = torch.softmax(pred, dim=1)
        # build the prediction dict for each output
        pred_dict = {'cls': pred, 'prob': prob, 'feat': features}

        return pred_dict

# Custom module to represent the residual using SVD components
class SVDResidualLinear(nn.Module):
    def __init__(self, in_features, out_features, r, bias=True, init_weight=None):
        super(SVDResidualLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.r = r  # Number of top singular values to exclude

        # Original weights (fixed)
        self.weight_main = nn.Parameter(torch.Tensor(out_features, in_features), requires_grad=False)
        if init_weight is not None:
            self.weight_main.data.copy_(init_weight)
        else:
            nn.init.kaiming_uniform_(self.weight_main, a=math.sqrt(5))

        # Bias
        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_features))
            nn.init.zeros_(self.bias)
        else:
            self.register_parameter('bias', None)

    def compute_current_weight(self):
        if self.S_residual is not None:
            return self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
        else:
            return self.weight_main

    def forward(self, x):
        if hasattr(self, 'U_residual') and hasattr(self, 'V_residual') and self.S_residual is not None:
            # Reconstruct the residual weight
            residual_weight = self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
            # Total weight is the fixed main weight plus the residual
            weight = self.weight_main + residual_weight
        else:
            # If residual components are not set, use only the main weight
            weight = self.weight_main

        return F.linear(x, weight, self.bias)

    def compute_orthogonal_loss(self):
        if self.S_residual is not None:
            # According to the properties of orthogonal matrices: A^TA = I
            UUT = torch.cat((self.U_r, self.U_residual), dim=1) @ torch.cat((self.U_r, self.U_residual), dim=1).t()
            VVT = torch.cat((self.V_r, self.V_residual), dim=0) @ torch.cat((self.V_r, self.V_residual), dim=0).t()
            # print(self.U_r.size(), self.U_residual.size())  # torch.Size([1024, 1023]) torch.Size([1024, 1])
            # print(self.V_r.size(), self.V_residual.size())  # torch.Size([1023, 1024]) torch.Size([1, 1024])
            # UUT = self.U_residual @ self.U_residual.t()
            # VVT = self.V_residual @ self.V_residual.t()

            # Construct an identity matrix
            UUT_identity = torch.eye(UUT.size(0), device=UUT.device)
            VVT_identity = torch.eye(VVT.size(0), device=VVT.device)

            # Using frobenius norm to compute loss
            loss = 0.5 * torch.norm(UUT - UUT_identity, p='fro') + 0.5 * torch.norm(VVT - VVT_identity, p='fro')
        else:
            loss = 0.0

        return loss

    def compute_keepsv_loss(self):
        if (self.S_residual is not None) and (self.weight_original_fnorm is not None):
            # Total current weight is the fixed main weight plus the residual
            weight_current = self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
            # Frobenius norm of current weight
            weight_current_fnorm = torch.norm(weight_current, p='fro')

            loss = torch.abs(weight_current_fnorm ** 2 - self.weight_original_fnorm ** 2)
            # loss = torch.abs(weight_current_fnorm ** 2 + 0.01 * self.weight_main_fnorm ** 2 - 1.01 * self.weight_original_fnorm ** 2)
        else:
            loss = 0.0

        return loss

    def compute_fn_loss(self):
        if (self.S_residual is not None):
            weight_current = self.weight_main + self.U_residual @ torch.diag(self.S_residual) @ self.V_residual
            weight_current_fnorm = torch.norm(weight_current, p='fro')

            loss = weight_current_fnorm ** 2
        else:
            loss = 0.0

        return loss


# Function to replace nn.Linear modules within self_attn modules with SVDResidualLinear
def apply_svd_residual_to_self_attn(model, r):
    for name, module in model.named_children():
        if 'self_attn' in name:
            # Replace nn.Linear layers in this module
            for sub_name, sub_module in module.named_modules():
                if isinstance(sub_module, nn.Linear):
                    # Get parent module within self_attn
                    parent_module = module
                    sub_module_names = sub_name.split('.')
                    for module_name in sub_module_names[:-1]:
                        parent_module = getattr(parent_module, module_name)
                    # Replace the nn.Linear layer with SVDResidualLinear
                    setattr(parent_module, sub_module_names[-1], replace_with_svd_residual(sub_module, r))
        else:
            # Recursively apply to child modules
            apply_svd_residual_to_self_attn(module, r)
    # After replacing, set requires_grad for residual components
    for param_name, param in model.named_parameters():
        if any(x in param_name for x in ['S_residual', 'U_residual', 'V_residual']):
            param.requires_grad = True
        else:
            param.requires_grad = False
    return model


# Function to replace a module with SVDResidualLinear
def replace_with_svd_residual(module, r):
    if isinstance(module, nn.Linear):
        in_features = module.in_features
        out_features = module.out_features
        bias = module.bias is not None

        # Create SVDResidualLinear module
        new_module = SVDResidualLinear(in_features, out_features, r, bias=bias, init_weight=module.weight.data.clone())

        if bias and module.bias is not None:
            new_module.bias.data.copy_(module.bias.data)

        new_module.weight_original_fnorm = torch.norm(module.weight.data, p='fro')

        # Perform SVD on the original weight
        U, S, Vh = torch.linalg.svd(module.weight.data, full_matrices=False)

        # Determine r based on the rank of the weight matrix
        r = min(r, len(S))  # Ensure r does not exceed the number of singular values

        # Keep top r singular components (main weight)
        U_r = U[:, :r]      # Shape: (out_features, r)
        S_r = S[:r]         # Shape: (r,)
        Vh_r = Vh[:r, :]    # Shape: (r, in_features)

        # Reconstruct the main weight (fixed)
        weight_main = U_r @ torch.diag(S_r) @ Vh_r

        # Calculate the frobenius norm of main weight
        new_module.weight_main_fnorm = torch.norm(weight_main.data, p='fro')

        # Set the main weight
        new_module.weight_main.data.copy_(weight_main)

        # Residual components (trainable)
        U_residual = U[:, r:]    # Shape: (out_features, n - r)
        S_residual = S[r:]       # Shape: (n - r,)
        Vh_residual = Vh[r:, :]  # Shape: (n - r, in_features)

        if len(S_residual) > 0:
            new_module.S_residual = nn.Parameter(S_residual.clone())
            new_module.U_residual = nn.Parameter(U_residual.clone())
            new_module.V_residual = nn.Parameter(Vh_residual.clone())

            new_module.S_r = nn.Parameter(S_r.clone(), requires_grad=False)
            new_module.U_r = nn.Parameter(U_r.clone(), requires_grad=False)
            new_module.V_r = nn.Parameter(Vh_r.clone(), requires_grad=False)
        else:
            new_module.S_residual = None
            new_module.U_residual = None
            new_module.V_residual = None

            new_module.S_r = None
            new_module.U_r = None
            new_module.V_r = None

        return new_module
    else:
        return module

'''
Training Params:

embeddings.class_embedding: False
embeddings.patch_embedding.weight: False
embeddings.position_embedding.weight: False
pre_layrnorm.weight: False
pre_layrnorm.bias: False

encoder.layers.0.self_attn.k_proj.weight_main: False
encoder.layers.0.self_attn.k_proj.bias: False
encoder.layers.0.self_attn.k_proj.S_residual: True
encoder.layers.0.self_attn.k_proj.U_residual: True
encoder.layers.0.self_attn.k_proj.V_residual: True
encoder.layers.0.self_attn.v_proj.weight_main: False
encoder.layers.0.self_attn.v_proj.bias: False
encoder.layers.0.self_attn.v_proj.S_residual: True
encoder.layers.0.self_attn.v_proj.U_residual: True
encoder.layers.0.self_attn.v_proj.V_residual: True
encoder.layers.0.self_attn.q_proj.weight_main: False
encoder.layers.0.self_attn.q_proj.bias: False
encoder.layers.0.self_attn.q_proj.S_residual: True
encoder.layers.0.self_attn.q_proj.U_residual: True
encoder.layers.0.self_attn.q_proj.V_residual: True
encoder.layers.0.self_attn.out_proj.weight_main: False
encoder.layers.0.self_attn.out_proj.bias: False
encoder.layers.0.self_attn.out_proj.S_residual: True
encoder.layers.0.self_attn.out_proj.U_residual: True
encoder.layers.0.self_attn.out_proj.V_residual: True
encoder.layers.0.layer_norm1.weight: False
encoder.layers.0.layer_norm1.bias: False
encoder.layers.0.mlp.fc1.weight: False
encoder.layers.0.mlp.fc1.bias: False
encoder.layers.0.mlp.fc2.weight: False
encoder.layers.0.mlp.fc2.bias: False
encoder.layers.0.layer_norm2.weight: False
encoder.layers.0.layer_norm2.bias: False

encoder.layers.1.self_attn.k_proj.weight_main: False
encoder.layers.1.self_attn.k_proj.bias: False
encoder.layers.1.self_attn.k_proj.S_residual: True
encoder.layers.1.self_attn.k_proj.U_residual: True
encoder.layers.1.self_attn.k_proj.V_residual: True
encoder.layers.1.self_attn.v_proj.weight_main: False
encoder.layers.1.self_attn.v_proj.bias: False
encoder.layers.1.self_attn.v_proj.S_residual: True
encoder.layers.1.self_attn.v_proj.U_residual: True
encoder.layers.1.self_attn.v_proj.V_residual: True
encoder.layers.1.self_attn.q_proj.weight_main: False
encoder.layers.1.self_attn.q_proj.bias: False
encoder.layers.1.self_attn.q_proj.S_residual: True
encoder.layers.1.self_attn.q_proj.U_residual: True
encoder.layers.1.self_attn.q_proj.V_residual: True
encoder.layers.1.self_attn.out_proj.weight_main: False
encoder.layers.1.self_attn.out_proj.bias: False
encoder.layers.1.self_attn.out_proj.S_residual: True
encoder.layers.1.self_attn.out_proj.U_residual: True
encoder.layers.1.self_attn.out_proj.V_residual: True
encoder.layers.1.layer_norm1.weight: False
encoder.layers.1.layer_norm1.bias: False
encoder.layers.1.mlp.fc1.weight: False
encoder.layers.1.mlp.fc1.bias: False
encoder.layers.1.mlp.fc2.weight: False
encoder.layers.1.mlp.fc2.bias: False
encoder.layers.1.layer_norm2.weight: False
encoder.layers.1.layer_norm2.bias: False

encoder.layers.2.self_attn.k_proj.weight_main: False
encoder.layers.2.self_attn.k_proj.bias: False
encoder.layers.2.self_attn.k_proj.S_residual: True
encoder.layers.2.self_attn.k_proj.U_residual: True
encoder.layers.2.self_attn.k_proj.V_residual: True
encoder.layers.2.self_attn.v_proj.weight_main: False
encoder.layers.2.self_attn.v_proj.bias: False
encoder.layers.2.self_attn.v_proj.S_residual: True
encoder.layers.2.self_attn.v_proj.U_residual: True
encoder.layers.2.self_attn.v_proj.V_residual: True
encoder.layers.2.self_attn.q_proj.weight_main: False
encoder.layers.2.self_attn.q_proj.bias: False
encoder.layers.2.self_attn.q_proj.S_residual: True
encoder.layers.2.self_attn.q_proj.U_residual: True
encoder.layers.2.self_attn.q_proj.V_residual: True
encoder.layers.2.self_attn.out_proj.weight_main: False
encoder.layers.2.self_attn.out_proj.bias: False
encoder.layers.2.self_attn.out_proj.S_residual: True
encoder.layers.2.self_attn.out_proj.U_residual: True
encoder.layers.2.self_attn.out_proj.V_residual: True
encoder.layers.2.layer_norm1.weight: False
encoder.layers.2.layer_norm1.bias: False
encoder.layers.2.mlp.fc1.weight: False
encoder.layers.2.mlp.fc1.bias: False
encoder.layers.2.mlp.fc2.weight: False
encoder.layers.2.mlp.fc2.bias: False
encoder.layers.2.layer_norm2.weight: False
encoder.layers.2.layer_norm2.bias: False
...
encoder.layers.23.self_attn.k_proj.weight_main: False
encoder.layers.23.self_attn.k_proj.bias: False
encoder.layers.23.self_attn.k_proj.S_residual: True
encoder.layers.23.self_attn.k_proj.U_residual: True
encoder.layers.23.self_attn.k_proj.V_residual: True
encoder.layers.23.self_attn.v_proj.weight_main: False
encoder.layers.23.self_attn.v_proj.bias: False
encoder.layers.23.self_attn.v_proj.S_residual: True
encoder.layers.23.self_attn.v_proj.U_residual: True
encoder.layers.23.self_attn.v_proj.V_residual: True
encoder.layers.23.self_attn.q_proj.weight_main: False
encoder.layers.23.self_attn.q_proj.bias: False
encoder.layers.23.self_attn.q_proj.S_residual: True
encoder.layers.23.self_attn.q_proj.U_residual: True
encoder.layers.23.self_attn.q_proj.V_residual: True
encoder.layers.23.self_attn.out_proj.weight_main: False
encoder.layers.23.self_attn.out_proj.bias: False
encoder.layers.23.self_attn.out_proj.S_residual: True
encoder.layers.23.self_attn.out_proj.U_residual: True
encoder.layers.23.self_attn.out_proj.V_residual: True
encoder.layers.23.layer_norm1.weight: False
encoder.layers.23.layer_norm1.bias: False
encoder.layers.23.mlp.fc1.weight: False
encoder.layers.23.mlp.fc1.bias: False
encoder.layers.23.mlp.fc2.weight: False
encoder.layers.23.mlp.fc2.bias: False
encoder.layers.23.layer_norm2.weight: False
encoder.layers.23.layer_norm2.bias: False

post_layernorm.weight: False
post_layernorm.bias: False
Number of total parameters: 303376480, tunable parameters: 196704


===> Load checkpoint done!
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 130/130 [01:07<00:00,  1.92it/s]
dataset: Celeb-DF-v2
acc: 0.7873882580333413
auc: 0.8674386218546616
eer: 0.21000704721634955
ap: 0.9322288761515111
pred: [0.9752515  0.6580601  0.75344455 ... 0.45359948 0.8914075  0.14674814]
video_auc: 0.9105750165234634
label: [1 1 0 ... 1 1 0]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 851/851 [07:17<00:00,  1.94it/s]
dataset: DeepFakeDetection
acc: 0.8606078424166698
auc: 0.9048725171315315
eer: 0.16390041493775934
ap: 0.9883843861944681
pred: [0.9912942  0.4690933  0.99789536 ... 0.8104649  0.9893     0.78386295]
video_auc: 0.9373875743738758
label: [1 1 1 ... 1 1 1]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 157/157 [01:21<00:00,  1.94it/s]
dataset: DFDCP
acc: 0.7002589125672177
auc: 0.8182703419711848
eer: 0.28125
ap: 0.9055071587990912
pred: [0.38877106 0.606897   0.5618232  ... 0.5063871  0.08330307 0.31897077]
video_auc: 0.849247887904389
label: [1 0 1 ... 1 0 0]
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1128/1128 [09:41<00:00,  1.94it/s]
dataset: DFDC
acc: 0.7389942337547128
auc: 0.8153807538695912
eer: 0.26452712297642716
ap: 0.8553135531715989
pred: [0.24556044 0.2040193  0.4257187  ... 0.82186127 0.9962172  0.50927925]
video_auc: 0.8395751948048426
label: [0 1 0 ... 0 1 0]
===> Test Done!
'''