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