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! '''