shunliwang

update

8bc3305 6 months ago

18.7 kB

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

	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