training/detectors/effort_patch_shuffle.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 albumentations as A
import loralib as lora
from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig

logger = logging.getLogger(__name__)


def get_clip_visual(model_name = "openai/clip-vit-base-patch16"):
    processor = AutoProcessor.from_pretrained(model_name)
    model = CLIPModel.from_pretrained(model_name)
    return processor, model.vision_model


def shuffle_patches(images: torch.Tensor, patch_size: int = 14) -> torch.Tensor:
    """
    Apply patch-level shuffling to the input images.
    images: [B, C, H, W]
    patch_size: patch size used by ViT (for example, 16)
    Returns: an image tensor with the same shape [B, C, H, W]
    """
    B, C, H, W = images.shape
    assert H % patch_size == 0 and W % patch_size == 0, \
        f"H ({H}) and W ({W}) must be divisible by patch_size ({patch_size})"

    num_patches_h = H // patch_size
    num_patches_w = W // patch_size
    num_patches = num_patches_h * num_patches_w

    # [B, C, H, W] -> [B, C, num_patches_h, patch_size, num_patches_w, patch_size]
    images = images.view(B, C, num_patches_h, patch_size, num_patches_w, patch_size)
    # -> [B, num_patches_h, num_patches_w, C, patch_size, patch_size]
    images = images.permute(0, 2, 4, 1, 3, 5).contiguous()
    # -> [B, num_patches, C, patch_size, patch_size]
    images = images.view(B, num_patches, C, patch_size, patch_size)

    # Shuffle patch order independently for each image.
    # permutation shape: [B, num_patches]
    perms = torch.stack(
        [torch.randperm(num_patches, device=images.device) for _ in range(B)],
        dim=0
    )
    # Use advanced indexing to perform the shuffle.
    batch_idx = torch.arange(B, device=images.device).unsqueeze(1).expand(B, num_patches)
    images = images[batch_idx, perms]  # [B, num_patches, C, patch_size, patch_size]

    # Restore the original image shape.
    images = images.view(B, num_patches_h, num_patches_w, C, patch_size, patch_size)
    # -> [B, C, num_patches_h, patch_size, num_patches_w, patch_size]
    images = images.permute(0, 3, 1, 4, 2, 5).contiguous()
    # -> [B, C, H, W]
    images = images.view(B, C, H, W)

    return images

def get_aug_transform():
    return A.Compose([
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
        A.HueSaturationValue(p=0.3),
        A.ImageCompression(quality_lower=40, quality_upper=100, p=0.1),
        A.GaussNoise(p=0.1),
        A.MotionBlur(p=0.1),
        A.CLAHE(p=0.1),
        A.ChannelShuffle(p=0.1),
        A.Cutout(p=0.1),
        A.RandomGamma(p=0.3),
        A.GlassBlur(p=0.3),
    ])


def data_aug(images: torch.Tensor) -> torch.Tensor:
    is_gpu = images.is_cuda
    aug = get_aug_transform()
    
    # Step 1: convert the batch tensor to batch numpy arrays (BHWC uint8 0-255).
    imgs_np = images.cpu().detach().numpy()
    imgs_np = np.transpose(imgs_np, (0, 2, 3, 1))  # BCHW -> BHWC
    imgs_np = (imgs_np * 255).astype(np.uint8)
    
    # Step 2: augment images one by one to avoid KeyError from batch-style arguments.
    imgs_aug_np = []
    for img in imgs_np:
        # Pass a single image with `image=img`, which is natively supported by Albumentations.
        aug_img = aug(image=img)["image"]
        imgs_aug_np.append(aug_img)
    imgs_aug_np = np.array(imgs_aug_np)  # convert back to batch numpy arrays
    
    # Step 3: convert back to a tensor while preserving the original logic.
    aug_tensor = torch.from_numpy(imgs_aug_np).permute(0, 3, 1, 2)
    aug_tensor = aug_tensor.float() / 255.0
    
    # Restore the original device.
    if is_gpu:
        aug_tensor = aug_tensor.cuda()
    
    return aug_tensor


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

    def build_clip_backbone(self,config):
        _, backbone = get_clip_visual(model_name=config['pretrained'])
        return backbone

    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-1)

        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])
        if self.training:
            #aug_image=data_aug(data_dict['image'])
            shuffle_images=shuffle_patches(data_dict['image'],14)
            feat = self.backbone(shuffle_images)['pooler_output']
            #feat1=self.backbone2(shuffle_images)['pooler_output']
        else:
            feat = self.backbone(data_dict['image'])['pooler_output']
            #feat1=self.backbone2(data_dict['image'])['pooler_output']
        # feat torch.Size([32, 1024])
        return feat#,feat1

    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)
        #features=features+f1
        #pred=pred+pred1
        # 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