import torch import torch.nn as nn import torch.distributed as dist import torch.nn.functional as F import torch.optim as optim from .kmeans_pytorch import kmeans from datamodules import RefTable from PIL import Image import os from utils import Dict class PrototypeLearner(nn.Module): def __init__(self, num_prototypes=64, prototype_dim=1024, lambda_reg=0.1, lr=0.01, iter_steps=50, verbose=True): super().__init__() self.num_prototypes = num_prototypes self.prototype_dim = prototype_dim self.lambda_reg = lambda_reg self.lr = lr self.iter_steps = iter_steps self.verbose = verbose def get_initial_prototypes(self, features, device): # features: [N_total, D] feats_norm = F.normalize(features, p=2, dim=-1) _, centers = kmeans( X=feats_norm, num_clusters=self.num_prototypes, distance='cosine', device=device, tqdm_flag=False ) return centers.to(device) # [K, D] def calculate_reconstruction_loss(self, targets, protos): protos_norm = F.normalize(protos, p=2, dim=-1) # Formula: W = T * P^T * (P * P^T + lambda * I)^(-1) # Ref: Eq. (2) in paper # P * P^T: Gram Matrix [K, K], (P * P^T + lambda * I)^(-1) p_gram = torch.matmul(protos_norm, protos_norm.t()) identity = torch.eye(self.num_prototypes, device=targets.device) inverse_term = torch.inverse(p_gram + self.lambda_reg * identity) # Mapping weights [N, K] # W = T * P^T * Inverse mapping_weights = torch.matmul(targets, protos_norm.t()) mapping_weights = torch.matmul(mapping_weights, inverse_term) # Reconstruct: T_hat = W * P reconstructed = torch.matmul(mapping_weights, protos_norm) return F.mse_loss(reconstructed, targets) def forward(self, features): """ features: [N_total, D] Return: Optimized Prototypes [K, D] """ device = features.device N, D = features.shape targets = F.normalize(features, p=2, dim=-1).detach() initial_protos = self.get_initial_prototypes(features, device) initial_protos_static = initial_protos.clone().detach() with torch.no_grad(): baseline_loss = self.calculate_reconstruction_loss(targets, initial_protos_static) if self.verbose: print(f"\n[ProtoLearner] Start. Baseline (K-Means) Reconstruction MSE: {baseline_loss.item():.6f}") prototypes = nn.Parameter(initial_protos.clone()) optimizer = optim.Adam([prototypes], lr=self.lr) for i in range(self.iter_steps): optimizer.zero_grad() loss = self.calculate_reconstruction_loss(targets, prototypes) loss.backward() optimizer.step() with torch.no_grad(): prototypes.data.copy_(F.normalize(prototypes.data, p=2, dim=-1)) if self.verbose and (i == 0 or (i + 1) % 10 == 0 or i == self.iter_steps - 1): with torch.no_grad(): curr_protos_norm = F.normalize(prototypes, p=2, dim=-1) init_protos_norm = F.normalize(initial_protos_static, p=2, dim=-1) shift_dist = torch.norm(curr_protos_norm - init_protos_norm, dim=-1).mean().item() cosine_sim = F.cosine_similarity(curr_protos_norm, init_protos_norm, dim=-1).mean().item() gain_pct = (baseline_loss.item() - loss.item()) / baseline_loss.item() * 100 print(f"Iter {i+1:02d}/{self.iter_steps} | " f"Loss: {loss.item():.6f} | " f"Gain: {gain_pct:.2f}% | " f"Shift(L2): {shift_dist:.4f} | " f"Sim(Cos): {cosine_sim:.4f}") return F.normalize(prototypes, p=2, dim=-1).detach() class PrototypeBank(nn.Module): def __init__(self, config, image_encoder, image_processor, num_prototypes=128, prototype_dim=1024): super().__init__() self.config = config self.num_prototypes = num_prototypes self.prototype_dim = prototype_dim self.image_patch_path = config.dataset.image_patch_path self.categories = config.dataset.categories[config.phase] self.aux = Dict(image_processor=image_processor, image_encoder=image_encoder) shape = (len(self.categories) + 1, num_prototypes, prototype_dim) self.register_buffer('prototypes', torch.zeros(shape)) self.register_buffer('prototype_flag', torch.tensor(False)) self.learner = PrototypeLearner( num_prototypes=num_prototypes, prototype_dim=prototype_dim, lambda_reg=0.1, lr=0.05, iter_steps=50 ) @property def image_processor(self): return self.aux.image_processor @property def image_encoder(self): return self.aux.image_encoder def build_prototypes(self): self.ref_table = RefTable()() self.cate_to_id = {cate: idx for (idx, cate) in enumerate(self.categories)} self.cate_to_id[''] = -1 if self.prototype_flag.item(): if (dist.is_initialized() and dist.get_rank() == 0) or (not dist.is_initialized()): print(f"[PrototypeBank] Prototypes loaded from checkpoint (Frozen). Skip calculation.") return is_dist = dist.is_initialized() rank = dist.get_rank() if is_dist else 0 device = next(iter(self.image_encoder.parameters())).device temp_prototypes = None if rank == 0: print("[PrototypeBank] Calculating prototypes online...") self.dense_features = {} for category in self.categories: files = list(self.ref_table[category].keys()) ref_images = [] for file in files: img = Image.open(os.path.join(self.image_patch_path, category, file)).convert('RGB') img = self.image_processor(images=img, return_tensors="pt", do_normalize=False)['pixel_values'].squeeze(0) ref_images.append(img) self.dense_features[category] = self.image_encoder(torch.stack(ref_images).to(device), mode='x_norm_patchtokens') prototype_list = [] for category in self.categories: flat_feats = self.dense_features[category].reshape(-1, self.prototype_dim) optimized_protos = self.learner(flat_feats) prototype_list.append(optimized_protos) prototype_list.append(torch.zeros_like(prototype_list[-1])) temp_prototypes = torch.stack(prototype_list) self.prototypes.copy_(temp_prototypes) self.prototype_flag.fill_(True) if is_dist: dist.broadcast(self.prototypes, src=0) dist.broadcast(self.prototype_flag, src=0) def get_prototypes(self, captions): return torch.stack(list(map(lambda caption: self.prototypes[self.cate_to_id[caption]], captions)))