training/detectors/clip_large_lsda.py

import os
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
from transformers import AutoProcessor, CLIPModel, ViTModel, ViTConfig
import loralib as lora
import copy

logger = logging.getLogger(__name__)

# Keep only intra-domain feature enhancement (centrifugal / affine / additive transforms, all enabled)
class LSDAAugmentor(nn.Module):
    def __init__(self, feature_dim=1024, beta_base=0.5, affine_theta_range=(-np.pi/6, np.pi/6), noise_ratio=0.01):
        super().__init__()
        self.feature_dim = feature_dim
        self.beta_base = beta_base  # base coefficient for the centrifugal transform
        self.affine_theta_range = affine_theta_range  # affine rotation angle range
        self.noise_ratio = noise_ratio  # ratio of noise to the feature standard deviation

    def centrifugal_trans(self, z: torch.Tensor) -> torch.Tensor:
        """Original version: adaptive centrifugal transform (the closer a sample is to the domain center, the farther it is pushed)"""
        batch_size = z.shape[0]
        mu = z.mean(dim=0, keepdim=True)  # domain center [1, 1024]
        dist = torch.norm(z - mu, dim=1, keepdim=True)  # distance from each sample to the center [B, 1]
        dist_max = dist.max() + 1e-6  # avoid division by zero
        beta = self.beta_base * (1 - dist / dist_max)  # the closer the distance, the larger β becomes
        z_aug = z + beta * (z - mu)  # apply the centrifugal transform
        return z_aug

    def affine_trans(self, z: torch.Tensor) -> torch.Tensor:
        """Affine transform that truly performs high-dimensional rotation (block-wise 2D rotations)"""
        batch_size, feat_dim = z.shape  # [32, 1024]
        device = z.device

        # 1. Validate the feature dimension: it must be even (1024 satisfies this requirement)
        assert feat_dim % 2 == 0, f"feature dimension{feat_dim}must be even in order to perform block-wise 2D rotations"
        num_blocks = feat_dim // 2  # 512 two-dimensional blocks

        # 2. Generate per-sample, per-block rotation angles (the actual rotation parameters)
        theta = torch.rand(batch_size, num_blocks, device=device) * (self.affine_theta_range[1] - self.affine_theta_range[0]) + self.affine_theta_range[0]
        # Generate per-sample scaling factors
        scale = torch.rand(batch_size, 1, device=device) * 0.1 + 0.95  # [32, 1]

        # 3. Compute the cos/sin values of the rotation matrix
        cos_theta = torch.cos(theta)  # [32, 512]
        sin_theta = torch.sin(theta)  # [32, 512]

        # 4. Split the 1024-dimensional features into 512 two-dimensional blocks
        z_reshaped = z.reshape(batch_size, num_blocks, 2)  # [32, 512, 2]

        # 5. Perform true 2D rotation on each two-dimensional block
        z_rotated = torch.stack([
            z_reshaped[..., 0] * cos_theta - z_reshaped[..., 1] * sin_theta,  # x' = xcosθ - ysinθ
            z_reshaped[..., 0] * sin_theta + z_reshaped[..., 1] * cos_theta   # y' = xsinθ + ycosθ
        ], dim=-1)  # [32, 512, 2]

        # 6. Reassemble by concatenation and apply global scaling
        z_aug = z_rotated.reshape(batch_size, feat_dim) * scale  # [32, 1024]

        return z_aug

    def additive_trans(self, z: torch.Tensor) -> torch.Tensor:
        """
    Apply adaptive additive noise augmentation to the input tensor based on the feature standard deviation

    Args:
        z (torch.Tensor): input feature tensor with shape [batch_size, feature_dim]

    Returns:
        torch.Tensor: augmented feature tensor with the same shape as the input

    Note:
        1. noise strength is jointly determined by the standard deviation of the feature dimension and the preset noise_ratio
        2. keep statistics computed independently for each batch dimension
        3. the generated noise follows a standard normal distribution N(0,1) and is then scaled
        """
        """Original version: adaptive additive noise (based on the feature standard deviation)"""
        # Compute the global feature standard deviation (batch-aware)
        z_std = z.std(dim=0, keepdim=True)
        noise = torch.randn_like(z) * z_std * self.noise_ratio
        z_aug = z + noise
        return z_aug

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        """Original core logic: sequentially stack all intra-domain augmentations during training, while validation/inference returns the original features"""
        if not self.training:
            return z

        # Step 1: centrifugal transform
        z = self.centrifugal_trans(z)
        # Step 2: affine transform
        z = self.affine_trans(z)
        # Step 3: additive noise
        z = self.additive_trans(z)

        return z


@DETECTOR.register_module(module_name='clip_large_fft_lsda')
class CLIP_Large_FFT_LSDA_Detector(AbstractDetector):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_classes = config['backbone_config']['num_classes']

        # Original components (kept intact)
        self.backbone = self.build_backbone(config)
        self.head = nn.Linear(1024, self.num_classes)
        self.loss_func = self.build_loss(config)

        # Keep only the intra-domain augmentor (no other new components)
        self.augmentor = LSDAAugmentor(feature_dim=1024)

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

    def build_loss(self, config):
        # Keep only the original classification loss, with no extra losses
        loss_class = LOSSFUNC[config['loss_func']]
        loss_func = loss_class()
        return loss_func

    def features(self, data_dict: dict) -> torch.tensor:
        """Core behavior: use augmented features for training and original features for validation/inference"""
        # Extract the original CLIP features
        feat = self.backbone(data_dict['image'])['pooler_output']  # [B, 1024]

        # During training, apply intra-domain augmentation; during validation/inference, return the original features directly
        feat = self.augmentor(feat)

        return feat

    def classifier(self, features: torch.tensor) -> torch.tensor:
        # Original classification head (unchanged)
        return self.head(features)

    def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
        # Keep only the original classification loss (no domain loss / distillation loss)
        label = data_dict['label']
        pred = pred_dict['cls']
        loss = self.loss_func(pred, label)
        loss_dict = {'overall': loss}
        return loss_dict

    def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
        # Original metric computation (unchanged)
        label = data_dict['label']
        pred = pred_dict['cls']
        acc, mAP = calculate_acc_for_train(label.detach(), pred.detach(), self.num_classes)
        metric_batch_dict = {'acc': acc, 'mAP': mAP}
        return metric_batch_dict

    def forward(self, data_dict: dict, inference=False) -> dict:
        """
        - When inference=True (validation/inference), the augmentor automatically returns the original features
        - When inference=False (training), the augmentor returns augmented features
        """
        # Control the augmentor mode (force eval during inference)
        if inference:
            self.augmentor.eval()
        else:
            self.augmentor.train()

        # Feature extraction (augmented during training, original during validation/inference)
        features = self.features(data_dict)
        # classification prediction (unchanged)
        pred = self.classifier(features)
        prob = torch.softmax(pred, dim=1)
        pred_dict = {'cls': pred, 'prob': prob, 'feat': features}
        return pred_dict


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