# Copyright (c) Alibaba, Inc. and its affiliates. import os.path as osp import random import json import os from glob import glob from tqdm import tqdm from contextlib import nullcontext import shutil import torch import torch.nn as nn from torch.nn import functional as F from torch.cuda.amp import GradScaler, autocast from dassl.engine import TRAINER_REGISTRY, TrainerX from dassl.metrics import compute_accuracy from dassl.utils import load_pretrained_weights, load_checkpoint from dassl.optim import build_optimizer, build_lr_scheduler from synthesis.feature_sample import IDFeatPool from ood.posthoc import applyReAct, applyBATS, applyASH from clip import clip from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer _tokenizer = _Tokenizer() perturb_methods = ['neg', 'zero', 'randn', 'randn_add', 'swap'] # def load_clip_to_cpu(cfg): backbone_name = cfg.MODEL.BACKBONE.NAME url = clip._MODELS[backbone_name] model_path = clip._download(url) try: # loading JIT archive model = torch.jit.load(model_path, map_location="cpu").eval() state_dict = None except RuntimeError: state_dict = torch.load(model_path, map_location="cpu") model = clip.build_model(state_dict or model.state_dict()) return model class TextEncoder(nn.Module): def __init__(self, clip_model): super().__init__() self.transformer = clip_model.transformer self.positional_embedding = clip_model.positional_embedding self.ln_final = clip_model.ln_final self.text_projection = clip_model.text_projection self.dtype = clip_model.dtype def forward(self, prompts, tokenized_prompts): ctx_len = prompts.shape[1] # TODO: compatible for dynamic context length x = prompts + self.positional_embedding.type(self.dtype)[:ctx_len] x = x.permute(1, 0, 2) # NLD -> LND x = self.transformer(x) x = x.permute(1, 0, 2) # LND -> NLD x = self.ln_final(x).type(self.dtype) # x.shape = [batch_size, n_ctx, transformer.width] # take features from the eot embedding (eot_token is the highest number in each sequence) x = x[torch.arange(x.shape[0]), tokenized_prompts.argmax(dim=-1)] @ self.text_projection return x class PromptLearner(nn.Module): def __init__(self, cfg, classnames, clip_model): super().__init__() n_cls = len(classnames) n_ctx = cfg.TRAINER.CATEX.N_CTX ctx_init = cfg.TRAINER.CATEX.CTX_INIT dtype = clip_model.dtype ctx_dim = clip_model.ln_final.weight.shape[0] clip_imsize = clip_model.visual.input_resolution cfg_imsize = cfg.INPUT.SIZE[0] assert cfg_imsize == clip_imsize, f"cfg_imsize ({cfg_imsize}) must equal to clip_imsize ({clip_imsize})" self.adjust_cls_promt = False self.cfg = cfg ctx_common = None if ctx_init and 'ensemble' not in ctx_init: # use given words to initialize context vectors ctx_init = ctx_init.replace("_", " ") n_ctx = len(ctx_init.split(" ")) prompt = clip.tokenize(ctx_init) with torch.no_grad(): embedding = clip_model.token_embedding(prompt).type(dtype) ctx_vectors = embedding[0, 1 : 1 + n_ctx, :] prompt_prefix = ctx_init else: # random initialization if cfg.TRAINER.CATEX.CSC: print("Initializing class-specific contexts") ctx_vectors = torch.empty(n_cls, n_ctx, ctx_dim, dtype=dtype) nn.init.normal_(ctx_vectors, std=0.02) else: print("Initializing a generic context") ctx_vectors = torch.empty(n_ctx, ctx_dim, dtype=dtype) nn.init.normal_(ctx_vectors, std=0.02) prompt_prefix = " ".join(["X"] * n_ctx) print(f'Initial context: "{prompt_prefix}"') print(f"Number of context words (tokens): {n_ctx}") self.ctx = nn.Parameter(ctx_vectors) # to be optimized self.ctx_cm = nn.Parameter(ctx_common) if ctx_common is not None else None if cfg.TRAINER.OOD_PROMPT: if cfg.TRAINER.OOD_PROMPT_NUM > 1: self.ctx_ood = [] for _ in range(cfg.TRAINER.OOD_PROMPT_NUM): ctx_ood = torch.empty(n_cls, n_ctx, ctx_dim, dtype=dtype) nn.init.normal_(ctx_ood, std=0.02) self.ctx_ood.append(nn.Parameter(ctx_ood)) self.ctx_ood = nn.ParameterList(self.ctx_ood) else: ## TODO: compatible for pre-trained weights ctx_ood = torch.empty(n_cls, n_ctx, ctx_dim, dtype=dtype) nn.init.normal_(ctx_ood, std=0.02) self.ctx_ood = nn.Parameter(ctx_ood) classnames = [name.replace("_", " ") for name in classnames] name_lens = [len(_tokenizer.encode(name)) for name in classnames] prompts = [prompt_prefix + " " + name + "." for name in classnames] tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts]) with torch.no_grad(): embedding = clip_model.token_embedding(tokenized_prompts).type(dtype) # These token vectors will be saved when in save_model(), # but they should be ignored in load_model() as we want to use # those computed using the current class names self.register_buffer("token_prefix", embedding[:, :1, :]) # SOS if self.adjust_cls_promt: self.token_suffix = nn.Parameter(embedding[:, 1 + n_ctx :, :]) else: self.register_buffer("token_suffix", embedding[:, 1 + n_ctx :, :]) # CLS, EOS self.n_cls = n_cls self.n_ctx = n_ctx self.tokenized_prompts = tokenized_prompts # torch.Tensor self.name_lens = name_lens self.class_token_position = cfg.TRAINER.CATEX.CLASS_TOKEN_POSITION def forward(self, perturb='none', ood_prompt=False, ood_prompt_idx=None): # ctx = self.ctx if ood_prompt: assert perturb == 'none', perturb if ood_prompt_idx is None: assert self.cfg.TRAINER.OOD_PROMPT_NUM == 1 ctx = self.ctx_ood else: ctx = self.ctx_ood[ood_prompt_idx] else: ctx = self.perturb_prompt(perturb) if ctx.dim() == 2: ctx = ctx.unsqueeze(0).expand(self.n_cls, -1, -1) if self.ctx_cm is not None: ctx_cm = self.ctx_cm.expand(self.n_cls, -1, -1) ctx = torch.cat((ctx_cm, ctx), dim=1) prefix = self.token_prefix suffix = self.token_suffix if self.class_token_position == "end": prompts = torch.cat( [ prefix, # (n_cls, 1, dim) ctx, # (n_cls, n_ctx, dim) suffix, # (n_cls, *, dim) ], dim=1, ) elif self.class_token_position == "middle": half_n_ctx = self.n_ctx // 2 prompts = [] for i in range(self.n_cls): name_len = self.name_lens[i] prefix_i = prefix[i : i + 1, :, :] class_i = suffix[i : i + 1, :name_len, :] suffix_i = suffix[i : i + 1, name_len:, :] ctx_i_half1 = ctx[i : i + 1, :half_n_ctx, :] ctx_i_half2 = ctx[i : i + 1, half_n_ctx:, :] prompt = torch.cat( [ prefix_i, # (1, 1, dim) ctx_i_half1, # (1, n_ctx//2, dim) class_i, # (1, name_len, dim) ctx_i_half2, # (1, n_ctx//2, dim) suffix_i, # (1, *, dim) ], dim=1, ) prompts.append(prompt) prompts = torch.cat(prompts, dim=0) elif self.class_token_position == "front": prompts = [] for i in range(self.n_cls): name_len = self.name_lens[i] prefix_i = prefix[i : i + 1, :, :] class_i = suffix[i : i + 1, :name_len, :] suffix_i = suffix[i : i + 1, name_len:, :] ctx_i = ctx[i : i + 1, :, :] prompt = torch.cat( [ prefix_i, # (1, 1, dim) class_i, # (1, name_len, dim) ctx_i, # (1, n_ctx, dim) suffix_i, # (1, *, dim) ], dim=1, ) prompts.append(prompt) prompts = torch.cat(prompts, dim=0) else: raise ValueError return prompts def perturb_prompt(self, method='none'): if method == 'none': return self.ctx coef_dict = { 'neg': [-1., 0.], 'zero': [0., 0.], 'randn': [0., 1.], 'randn_add': [1., 1.], 'swap': [0., 1.] } assert method in coef_dict ctx = self.ctx if ctx.dim() == 2: ctx = ctx.unsqueeze(0).expand(self.n_cls, -1, -1) ncls, nctx, ndim = ctx.shape assert nctx > 1 for i in range(self.cfg.TRAINER.ID_PERTUR_NUM): # perturb one prompt for each class ctx_ind = torch.randint(0, nctx, size=(ncls,)) cls_ind = torch.arange(ncls) src_mask = torch.ones((ncls, nctx, 1)).type_as(ctx) src_mask[cls_ind, ctx_ind] = 0. src_ctx = ctx[cls_ind, ctx_ind].detach() if method == 'swap': ori_ind = torch.arange(ncls) while True: rand_ind = torch.randperm(ncls) if (ori_ind != rand_ind).all(): noise = src_ctx[rand_ind] break else: noise = torch.randn_like(ctx[:, 0, :]) src_coef, noise_coef = coef_dict[method] perturb = torch.zeros_like(ctx) perturb[cls_ind, ctx_ind] = src_coef * src_ctx + noise_coef * noise ctx = ctx * src_mask + perturb * (1. - src_mask) return ctx class CustomCLIP(nn.Module): def __init__(self, cfg, classnames, clip_model): super().__init__() self.classnames = classnames self.token_embedding = clip_model.token_embedding self.prompt_learner = PromptLearner(cfg, classnames, clip_model) self.tokenized_prompts = self.prompt_learner.tokenized_prompts self.image_encoder = clip_model.visual self.text_encoder = TextEncoder(clip_model) self.logit_scale = clip_model.logit_scale if not cfg.TRAINER.OOD_TEST else torch.zeros_like(clip_model.logit_scale) self.dtype = clip_model.dtype self.feat_dim = clip_model.text_projection.data.shape[1] self.text_feature_ensemble = self.prompt_ensemble() if cfg.TRAINER.CATEX.CTX_INIT == 'ensemble' else None @torch.no_grad() def prompt_ensemble(self, learned_text_features=None): if learned_text_features is None: imagenet_templates = [ # for NPOS 'a photo of a {}.', 'a blurry photo of a {}.', 'a black and white photo of a {}.', 'a low contrast photo of a {}.', 'a high contrast photo of a {}.', 'a bad photo of a {}.', 'a good photo of a {}.', 'a photo of a small {}.', 'a photo of a big {}.', 'a photo of the {}.', 'a blurry photo of the {}.', 'a black and white photo of the {}.', 'a low contrast photo of the {}.', 'a high contrast photo of the {}.', 'a bad photo of the {}.', 'a good photo of the {}.', 'a photo of the small {}.', 'a photo of the big {}.', ] else: imagenet_templates = [ # for MCM 'a photo of a {}.', 'a blurry photo of a {}.', 'a photo of many {}.', 'a black and white photo of a {}.', 'a photo of the large {}.', 'a photo of the small {}.', ] lambd = 0.5 dtype = self.text_encoder.dtype self.text_encoder = self.text_encoder.cuda() self.token_embedding = self.token_embedding.cuda() text_feature = [] for ci, classname in enumerate(self.classnames): texts = [template.format(classname) for template in imagenet_templates] # format with class texts = clip.tokenize(texts).cuda() # tokenize embedding = self.token_embedding(texts).type(dtype) class_embeddings = self.text_encoder(embedding, texts) # embed with text encoder class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True) if learned_text_features is not None: class_embeddings = torch.cat((class_embeddings, lambd * learned_text_features[ci:ci+1])) class_embedding = class_embeddings.mean(dim=0) class_embedding = class_embedding / class_embedding.norm() text_feature.append(class_embedding) text_feature = torch.stack(text_feature, dim=0).type(dtype) return text_feature def get_text_features(self, perturb='none', ood_prompt=False, ood_prompt_idx=None, return_norm=True): if self.text_feature_ensemble is not None and ood_prompt is False and perturb == 'none': assert return_norm text_features = self.text_feature_ensemble else: prompts = self.prompt_learner(perturb, ood_prompt=ood_prompt, ood_prompt_idx=ood_prompt_idx) tokenized_prompts = self.tokenized_prompts text_features = self.text_encoder(prompts, tokenized_prompts) if return_norm: text_features = text_features / text_features.norm(dim=-1, keepdim=True) return text_features def get_all_ood_text_features(self): x = [] for ood_prompt_idx in range(self.prompt_learner.cfg.TRAINER.OOD_PROMPT_NUM): x.append(self.get_text_features(ood_prompt=True, ood_prompt_idx=ood_prompt_idx)) return torch.stack(x, dim=1) # shape(1000,5,512) def get_logits(self, image_features, text_features, logit_scale=None): if logit_scale is None: logit_scale = self.logit_scale.exp() if text_features.dim() == 2: logits = image_features.float() @ text_features.float().t() else: n = text_features.size(0) logits = torch.bmm(image_features.unsqueeze(0).repeat(n,1,1), text_features.transpose(1,2)).max(dim=0)[0] #.mean(dim=0) return logits * logit_scale def forward(self, image, perturb='none', ood_prompt=False, return_feat=False, return_norm=True, posthoc=None): if len(image.shape) == 2: image_features = image else: assert len(image.shape) == 4 image_features = self.image_encoder(image.type(self.dtype)) if return_feat and not return_norm: ret_feat = image_features.detach().clone() if posthoc == 'apply_react': image_features = applyReAct(image_features) elif posthoc == 'apply_bats': image_features = applyBATS(image_features) elif posthoc == 'apply_ash': image_features = applyASH(image_features) image_features = image_features / image_features.norm(dim=-1, keepdim=True) if return_feat and return_norm: ret_feat = image_features.detach().clone() text_features = self.get_text_features(perturb, ood_prompt=ood_prompt) logits = self.get_logits(image_features, text_features) if return_feat: return ret_feat, text_features, logits else: return logits def calc_prompt_loss(self, clean_feat=None, perturb_feat=None): l_intra, l_inter = 0., 0. # # 1. class-specific prompts should not be similar to class-name prompts # csc_prompts = F.normalize(self.prompt_learner.ctx, p=2, dim=-1) # cls_prompts = F.normalize(self.prompt_learner.token_suffix[:, :max(self.prompt_learner.name_lens), :], p=2, dim=-1) # prompts = torch.cat((csc_prompts, cls_prompts), dim=1) # shape(ncls, nctx, ndim) # similarity = torch.bmm(prompts, prompts.transpose(1,2)) # shape(ncls, nctx, nctx) # diag = torch.arange(similarity.shape[1]) # similarity[:, diag, diag] = -1. # l_intra += similarity.max(dim=-1)[0].relu().mean() # 2. prompts should obviously affect the text-feature if clean_feat is None: clean_feat = self.get_text_features() if perturb_feat is None: # with torch.no_grad(): perturb_feat = self.get_text_features(random.choice(perturb_methods)) similarity = (perturb_feat * clean_feat).sum(dim=1) l_inter += (similarity - 0.8).relu().mean() return l_intra + l_inter def calc_ood_prompt_loss(self, image, logits, label): perturb_logits = self.forward(image, perturb=random.choice(perturb_methods)) bi = torch.arange(image.shape[0]) intra_loss = (perturb_logits[bi, label] - logits[bi, label]).relu().mean() inter_loss = -(perturb_logits.mean(1) - torch.logsumexp(perturb_logits, dim=1)).mean() return intra_loss + inter_loss @TRAINER_REGISTRY.register() class CATEX(TrainerX): """Context Optimization (CATEX). Learning to Prompt for Vision-Language Models https://arxiv.org/abs/2109.01134 """ def __init__(self, cfg): super().__init__(cfg) if cfg.TRAINER.OOD_TRAIN or self.is_large_ID(): if self.is_large_ID(): nsample = 1200 # 1200 self.id_pool = IDFeatPool(self.model.prompt_learner.n_cls, nsample, self.model.feat_dim, mode='npos', device='cuda:0') if cfg.TRAINER.ID_FEAT_PRELOAD != '': queue = torch.load(cfg.TRAINER.ID_FEAT_PRELOAD).to(self.id_pool.queue.device) self.id_pool.queue = queue[:, :nsample, :] self.id_pool.class_ptr += nsample else: from torch.utils.data import DataLoader, Subset from ood.datasets import TinyImages, InfiniteDataLoader assert 'cifar' in self.dm.dataset.dataset_name data_root = osp.abspath(osp.expanduser(cfg.DATASET.ROOT)) ood_set = TinyImages(data_root, transform=self.train_loader_x.dataset.transform) self.ood_loader = InfiniteDataLoader(ood_set, batch_size=cfg.DATALOADER.TRAIN_X.BATCH_SIZE, shuffle=False, num_workers=self.train_loader_x.num_workers, pin_memory=True) # drop_last=True, from ood.losses import LogitNormLoss self.ce_criterion = LogitNormLoss() if cfg.TRAINER.LOGIT_NORM else nn.CrossEntropyLoss() def is_large_ID(self): # return True return 'imagenet' in self.dm.dataset.dataset_name def check_cfg(self, cfg): assert cfg.TRAINER.CATEX.PREC in ["fp16", "fp32", "amp"] def build_model(self): cfg = self.cfg classnames = self.dm.dataset.classnames print(f"Loading CLIP (backbone: {cfg.MODEL.BACKBONE.NAME})") clip_model = load_clip_to_cpu(cfg) if cfg.TRAINER.CATEX.PREC == "fp32" or cfg.TRAINER.CATEX.PREC == "amp": # CLIP's default precision is fp16 clip_model.float() print("Building custom CLIP") self.model = CustomCLIP(cfg, classnames, clip_model) print("Turning off gradients in both the image and the text encoder") for name, param in self.model.named_parameters(): if "prompt_learner" not in name: param.requires_grad_(False) if cfg.MODEL.INIT_WEIGHTS: load_pretrained_weights(self.model.prompt_learner, cfg.MODEL.INIT_WEIGHTS) self.model.to(self.device) # NOTE: only give prompt_learner/image_encoder to the optimizer self.optim = build_optimizer(self.model.prompt_learner, cfg.OPTIM) self.sched = build_lr_scheduler(self.optim, cfg.OPTIM) self.register_model("prompt_learner", self.model.prompt_learner, self.optim, self.sched) self.scaler = GradScaler() if cfg.TRAINER.CATEX.PREC == "amp" else None # # Note that multi-gpu training could be slow because CLIP's size is # # big, which slows down the copy operation in DataParallel # device_count = torch.cuda.device_count() # if device_count > 1: # print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!") # self.model = nn.DataParallel(self.model) def build_data_loader(self): """Create essential data-related attributes. A re-implementation of this method must create the same attributes (self.dm is optional). """ super().build_data_loader() if self.cfg.TRAINER.FEAT_AS_INPUT and not self.cfg.TRAINER.OOD_TEST: from ood.datasets import CLIPFeatDataset self.load_shuffle = False self.train_loader_x = torch.utils.data.DataLoader( CLIPFeatDataset(self.dm.dataset.dataset_dir+'/clip_feat', self.start_epoch), batch_size=self.dm.train_loader_x.batch_size, shuffle=self.load_shuffle, num_workers=self.dm.train_loader_x.num_workers, pin_memory=True, drop_last=False, ) def calc_loss(self, logits, label, image=None, image_features=None, text_features=None, return_norm=False): nb, ncls = logits.shape # 1. classification if self.cfg.TRAINER.OOD_PROMPT and self.epoch >= self.cfg.TRAINER.START_EPOCH \ and self.is_large_ID() and self.id_pool.ready(): if not return_norm: image_features = F.normalize(image_features, p=2, dim=1) ood_text_features = self.model.get_text_features(ood_prompt=True) if self.cfg.TRAINER.OOD_PROMPT_CE_LOSS: logits = self.model.get_logits(image_features, torch.cat((text_features, ood_text_features))) logits[torch.arange(nb), label+ncls] = -10. # generally -inf loss = 2. * self.ce_criterion(logits, label) # 2. prompt perturbation perturbed_text_features = None if self.cfg.TRAINER.OOD_PROMPT and self.cfg.TRAINER.ID_PERTURB_LOSS and self.epoch >= self.cfg.TRAINER.START_EPOCH: with torch.no_grad(): perturbed_text_features = self.model.get_text_features(perturb=random.choice(perturb_methods)) loss += 0.1 * self.model.calc_prompt_loss(text_features, perturbed_text_features) # 3. outlier exposure assert text_features is not None if self.is_large_ID(): if self.id_pool.ready() and self.cfg.TRAINER.OOD_PROMPT and self.epoch >= self.cfg.TRAINER.START_EPOCH: if logits.size(0) < self.id_pool.queue.size(0): cls_mask = torch.unique(label).cpu() else: cls_mask = None if self.cfg.TRAINER.OOD_ANCHOR: if self.cfg.TRAINER.ID_PERTURB_LOSS and False: perturbed_text_features = self.model.get_text_features(perturb=random.choice(perturb_methods)) logit_scale = self.model.logit_scale.exp() id_pos_sim = (image_features * text_features[label]).sum(dim=-1) * logit_scale id_neg_sim = (image_features * perturbed_text_features[label]).sum(dim=-1) * logit_scale loss += F.cross_entropy(torch.stack((id_pos_sim, id_neg_sim), dim=1), torch.zeros((len(id_pos_sim),), dtype=torch.long, device=self.device)) * 0.5 elif perturbed_text_features is None: with torch.no_grad(): perturbed_text_features = self.model.get_text_features(perturb=random.choice(perturb_methods)) text_anchors = torch.stack((text_features, perturbed_text_features), dim=1).detach() else: text_anchors = None ood_features, ood_labels = self.id_pool.gen_ood(anchors=text_anchors, device=self.device, cls_mask=cls_mask) if self.cfg.TRAINER.OOD_OE_LOSS: ood_logits = self.model.get_logits(ood_features, text_features, logit_scale=1.) loss += 0.5 * -(ood_logits.mean(1) - torch.logsumexp(ood_logits, dim=1)).mean() if self.cfg.TRAINER.OOD_PROMPT: # ood_text_features = self.model.get_text_features(ood_prompt=True) if self.cfg.TRAINER.OOD_PROMPT_ORTH: assert self.cfg.TRAINER.OOD_PROMPT_NUM > 1 all_ood_text_features = self.model.get_all_ood_text_features() # (1000,5,512) x (1000,512,5) -> (1000,5,5) ood_sim_matrix = torch.bmm(all_ood_text_features, all_ood_text_features.transpose(1,2)) ood_text_num = ood_sim_matrix.shape[-1] zrange = torch.arange(ood_text_num) ood_sim_matrix[:, zrange, zrange] = 0. loss += 0.1 * ood_sim_matrix.mean() if self.cfg.TRAINER.OOD_PROMPT_CE_LOSS: ood_logits = self.model.get_logits(ood_features, torch.cat((ood_text_features, text_features))) ood_logits[torch.arange(ood_logits.shape[0]), ood_labels+ncls] = -10. # generally -inf loss += 0.5 * self.ce_criterion(ood_logits, ood_labels) if self.cfg.TRAINER.OOD_PROMPT_MARGIN_LOSS: if self.cfg.TRAINER.OOD_PROMPT_MARGIN_SOFT_LOSS: logit_scale = self.model.logit_scale.exp() else: logit_scale = 1. id_pos_sim = (image_features * text_features[label]).sum(dim=-1) * logit_scale id_neg_sim = (image_features * ood_text_features[label]).sum(dim=-1) * logit_scale ood_pos_sim = (ood_features * ood_text_features[ood_labels]).sum(dim=-1) * logit_scale ood_neg_sim = (ood_features * text_features[ood_labels]).sum(dim=-1) * logit_scale # id_pos_sim = (image_features @ text_features.T).max(dim=-1)[0] * logit_scale # id_neg_sim = (image_features @ ood_text_features.T).max(dim=-1)[0] * logit_scale # ood_pos_sim = (ood_features @ ood_text_features.T).max(dim=-1)[0] * logit_scale # ood_neg_sim = (ood_features @ text_features.T).max(dim=-1)[0] * logit_scale if self.cfg.TRAINER.OOD_PROMPT_MARGIN_SOFT_LOSS: loss += F.cross_entropy(torch.stack((id_pos_sim, id_neg_sim), dim=1), torch.zeros((len(id_pos_sim),), dtype=torch.long, device=self.device)) + \ F.cross_entropy(torch.stack((ood_pos_sim, ood_neg_sim), dim=1), torch.zeros((len(ood_pos_sim),), dtype=torch.long, device=self.device)) else: loss += (id_neg_sim - id_pos_sim).relu().mean() + (ood_neg_sim - ood_pos_sim).relu().mean() else: ood_data, _ = next(self.ood_loader.__iter__()) ood_data = ood_data.to(self.device) ood_logits = self.model(ood_data) #/ self.model.logit_scale.exp() loss += 0.1 * -(ood_logits.mean(1) - torch.logsumexp(ood_logits, dim=1)).mean() return loss def forward_backward(self, batch): image, label = self.parse_batch_train(batch) prec = self.cfg.TRAINER.CATEX.PREC if prec == "amp": with autocast(): output = self.model(image) loss = F.cross_entropy(output, label) self.optim.zero_grad() self.scaler.scale(loss).backward() self.scaler.step(self.optim) self.scaler.update() else: output = self.model(image) loss = F.cross_entropy(output, label) self.model_backward_and_update(loss) loss_summary = { "loss": loss.item(), "acc": compute_accuracy(output, label)[0].item(), } if (self.batch_idx + 1) == self.num_batches: self.update_lr() return loss_summary def forward_backward_ood(self, batch): image, label = self.parse_batch_train(batch) prec = self.cfg.TRAINER.CATEX.PREC return_norm = False if prec == "amp": with autocast(): img_feat, text_feat, output = \ self.model(image, return_feat=True, return_norm=return_norm) self.id_pool.update(img_feat.detach(), label) loss = self.calc_loss(output, label, image, img_feat, text_feat, return_norm=return_norm) self.optim.zero_grad() self.scaler.scale(loss).backward() self.scaler.step(self.optim) self.scaler.update() else: img_feat, text_feat, output = \ self.model(image, return_feat=True, return_norm=return_norm) self.id_pool.update(img_feat.detach(), label) loss = self.calc_loss(output, label, image, img_feat, text_feat, return_norm=return_norm) self.model_backward_and_update(loss) loss_summary = { "loss": loss.item(), "acc": compute_accuracy(output, label)[0].item(), } if (self.batch_idx + 1) == self.num_batches: self.update_lr() return loss_summary def parse_batch_train(self, batch): input = batch["img"] label = batch["label"] input = input.to(self.device) label = label.to(self.device) return input, label def load_model(self, directory, epoch=None): if not directory: print("Note that load_model() is skipped as no pretrained model is given") return names = self.get_model_names() # By default, the best model is loaded model_file = "model-best.pth.tar" if epoch is not None: model_file = "model.pth.tar-" + str(epoch) for name in names: model_path = osp.join(directory, name, model_file) if not osp.exists(model_path): raise FileNotFoundError('Model not found at "{}"'.format(model_path)) checkpoint = load_checkpoint(model_path) state_dict = checkpoint["state_dict"] epoch = checkpoint["epoch"] # Ignore fixed token vectors model_dict = self._models[name].state_dict() if "token_prefix" in state_dict: del state_dict["token_prefix"] if "token_suffix" in state_dict: del state_dict["token_suffix"] assert all(k in model_dict for k in state_dict) print("Loading weights to {} {} " 'from "{}" (epoch = {})'.format(name, list(state_dict.keys()), model_path, epoch)) # set strict=False self._models[name].load_state_dict(state_dict, strict=False) if self.cfg.TRAINER.CATEX.CTX_INIT: assert self.cfg.TRAINER.CATEX.CTX_INIT == 'ensemble_learned' text_feature = self.model.get_text_features() self.model.text_feature_ensemble = self.model.prompt_ensemble(text_feature) def load_model_vanilla(self, directory, name, epoch=None): if not directory: print("Note that load_model() is skipped as no pretrained model is given") return # By default, the best model is loaded model_file = "model-best.pth.tar" if epoch is not None: model_file = "model.pth.tar-" + str(epoch) model_path = osp.join(directory, name, model_file) if not osp.exists(model_path): raise FileNotFoundError('Model not found at "{}"'.format(model_path)) checkpoint = load_checkpoint(model_path) state_dict = checkpoint["state_dict"] epoch = checkpoint["epoch"] print("Loading vanilla weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch)) # set strict=False model = self.model.module if torch.cuda.device_count() > 1 else self.model getattr(model, name).load_state_dict(state_dict, strict=False) def before_epoch(self): if self.cfg.TRAINER.FEAT_AS_INPUT: if not self.load_shuffle: self.train_loader_x.dataset.load_data(self.epoch) @torch.no_grad() def test_ood(self, split=None, model_directory=''): """A generic OOD testing pipeline.""" from tqdm import tqdm import os import os.path as osp from torch.utils.data import DataLoader import numpy as np from ood.datasets import CLIPFeatDataset from ood.datasets import SCOODDataset, LargeOODDataset, SemanticOODDataset, ClassOODDataset from ood.metrics import get_msp_scores, get_measures self.set_model_mode("eval") self.evaluator.reset() if split is None: split = self.cfg.TEST.SPLIT if self.cfg.TRAINER.FEAT_AS_INPUT: feat_data_dir = self.dm.dataset.dataset_dir+'/clip_feat' if not osp.exists(feat_data_dir): self.cache_feat(split=split, is_ood=False) data_loader = DataLoader( CLIPFeatDataset(feat_data_dir, self.start_epoch, split='test'), batch_size=self.test_loader.batch_size, shuffle=False, num_workers=self.test_loader.num_workers, pin_memory=True, drop_last=False, ) else: if split == "val" and self.val_loader is not None: data_loader = self.val_loader else: split = "test" # in case val_loader is None data_loader = self.test_loader lab2cname = self.dm.dataset.lab2cname ood_cfg = { 'SCOOD': ['texture', 'svhn', 'cifar', 'tin', 'lsun', 'places365'], 'LargeOOD': ['inaturalist', 'sun', 'places', 'texture'], } data_root = osp.abspath(osp.expanduser(self.cfg.DATASET.ROOT)) ood_type = 'SCOOD' if 'cifar' in self.dm.dataset.dataset_name else 'LargeOOD' # LargeOOD, ClassOOD if 'apply_' in self.cfg.TRAINER.OOD_INFER_OPTION: posthoc = self.cfg.TRAINER.OOD_INFER_OPTION else: posthoc = None if self.cfg.TRAINER.OOD_PROMPT: if self.cfg.TRAINER.OOD_PROMPT_NUM > 1: ood_text_features = torch.stack([self.model.get_text_features(ood_prompt=True, ood_prompt_idx=i) for i in range(self.cfg.TRAINER.OOD_PROMPT_NUM)]) else: ood_text_features = self.model.get_text_features(ood_prompt=True) if self.cfg.TRAINER.CATEX.CTX_INIT: assert self.cfg.TRAINER.CATEX.CTX_INIT == 'ensemble_learned' ood_text_features = self.model.prompt_ensemble(ood_text_features) min_thresh = 0.51 if any(flag in model_directory for flag in ['/imagenet/', '/imagenet100-MCM-SCTX8-Orth/']) else 0.5 self.model.text_feature_ensemble = self.model.get_text_features() print(f"Evaluate on the *{split}* set") if self.cfg.TRAINER.OOD_INFER_OPTION == 'save_res': save_dir = f'{model_directory}/restore' os.makedirs(save_dir, exist_ok=True) with open(f'{save_dir}/lab2cname.json', 'w+') as f: json.dump(lab2cname, f, indent=4) text_features = self.model.get_text_features() torch.save(text_features.cpu(), f'{save_dir}/in_text_features.pt') if self.cfg.TRAINER.OOD_PROMPT: torch.save(ood_text_features.cpu(), f'{save_dir}/ood_text_features.pt') im_feats, im_labels = [], [] if self.cfg.TRAINER.OOD_INFER_OPTION == 'resume_res': resume_dir = 'weights/imagenet100-MCM/CATEX/vit_b16_ep50_-1shots/nctx16_cscTrue_ctpend/seed1/restore' resume_image_features = torch.load(f'{resume_dir}/in_image_features.pt').to(self.device) resume_image_labels = torch.load(f'{resume_dir}/in_labels.pt').to(self.device) resume_text_features = torch.load(f'{resume_dir}/in_text_features.pt').to(self.device) if self.cfg.TRAINER.OOD_PROMPT: resume_ood_text_features = torch.load(f'{resume_dir}/ood_text_features.pt').to(self.device) with open(f'{resume_dir}/lab2cname.json', 'r') as f: resume_lab2cname = json.load(f) resume_lab2cname = {int(k): v for k, v in resume_lab2cname.items()} label_offset = resume_image_labels.max().item() + 1 resume_image_labels += label_offset text_features = self.model.get_text_features() merged_text_features = torch.cat((text_features, resume_text_features), dim=0) if self.cfg.TRAINER.OOD_PROMPT: # TODO: not implemented merged_ood_text_features = torch.cat((ood_text_features, resume_ood_text_features), dim=0) score_list = [] base_acc, novel_acc = [], [] near_ood_flag = [] all_logits = [] for batch_idx, batch in enumerate(tqdm(data_loader)): input, label = self.parse_batch_test(batch) image_features, _, output = \ self.model(input, return_feat=True, return_norm=True, posthoc=posthoc) if self.cfg.TRAINER.OOD_PROMPT: ood_logits = self.model.get_logits(image_features, ood_text_features, logit_scale=1.) # all_logits.append(torch.stack((output, ood_logits), dim=1)) if self.cfg.TRAINER.OOD_INFER_INTEGRATE: id_score = F.softmax(torch.stack((output, ood_logits), dim=1), dim=1)[:, 0, :] output *= id_score.clamp(min=min_thresh) else: ood_logits = None if self.cfg.TRAINER.OOD_INFER_OPTION == 'save_res': im_feats.append(image_features.cpu()) im_labels.append(label.cpu()) if self.cfg.TRAINER.OOD_INFER_OPTION == 'resume_res': start = data_loader.batch_size * batch_idx end = start + input.shape[0] merged_image_features = torch.cat((image_features, resume_image_features[start:end]), dim=0) label = torch.cat((label, resume_image_labels[start:end]), dim=0) output = merged_image_features @ merged_text_features.t() acc = output.argmax(dim=1) == label base_acc.append(acc[input.shape[0]:].cpu()) novel_acc.append(acc[:input.shape[0]].cpu()) if hasattr(self.dm.dataset, 'valid_classes'): # if self.cfg.TRAINER.OOD_PROMPT: # raise NotImplementedError output[:, ~self.dm.dataset.valid_classes] = -1. scores = get_msp_scores(output[:, self.dm.dataset.valid_classes]) else: scores, ood_flag = get_msp_scores(output, ood_logits, self.cfg.TRAINER.OOD_INFER, ret_near_ood=True) near_ood_flag.append(ood_flag) score_list.append(scores.detach().cpu().numpy()) self.evaluator.process(output, label) in_scores = np.concatenate(score_list, axis=0) results = self.evaluator.evaluate() if self.cfg.TRAINER.OOD_PROMPT and len(near_ood_flag) and near_ood_flag[0] is not None: print('NearOOD FPR:', torch.cat(near_ood_flag).sum().item() / len(in_scores)) if self.cfg.TRAINER.OOD_INFER_OPTION == 'save_res': torch.save(torch.cat(im_feats), f'{save_dir}/in_image_features.pt') torch.save(torch.cat(im_labels), f'{save_dir}/in_labels.pt') if len(all_logits): torch.save(torch.cat(all_logits), f'{save_dir}/in_logits_all.pt') if self.cfg.TRAINER.OOD_INFER_OPTION == 'resume_res': print(f'Base: {torch.cat(base_acc).float().mean(): .4f}. Novel: {torch.cat(novel_acc).float().mean(): .4f}') auroc_list, aupr_list, fpr95_list = [], [], [] ood_tpr_list = [] save_lines = [] for ood_name in ood_cfg[ood_type]: ood_set = eval(f'{ood_type}Dataset')(osp.join(data_root, ood_type), id_name=self.dm.dataset.dataset_name, ood_name=ood_name, transform=self.test_loader.dataset.transform) if self.cfg.TRAINER.FEAT_AS_INPUT: feat_data_dir = f'{data_root}/{ood_type}/clip_feat/{ood_name}' if not osp.exists(feat_data_dir): self.cache_feat(split='test', is_ood=True) ood_loader = DataLoader( CLIPFeatDataset(feat_data_dir, epoch=None, split='test'), batch_size=self.cfg.DATALOADER.TEST.BATCH_SIZE, shuffle=False, num_workers=data_loader.num_workers, pin_memory=True, drop_last=False, ) else: ood_loader = DataLoader(ood_set, batch_size=self.cfg.DATALOADER.TEST.BATCH_SIZE, shuffle=False, num_workers=data_loader.num_workers, drop_last=False, pin_memory=True) ood_score_list, sc_labels_list, ood_pred_list = [], [], [] near_ood_flag = [] all_logits = [] for batch_idx, batch in enumerate(tqdm(ood_loader)): if self.cfg.TRAINER.FEAT_AS_INPUT: images, sc_labels = self.parse_batch_test(batch) else: images, sc_labels = batch images = images.to(self.device) image_features, _, output = \ self.model(images, return_feat=True, return_norm=True, posthoc=posthoc) if self.cfg.TRAINER.OOD_PROMPT: ood_logits = self.model.get_logits(image_features, ood_text_features, logit_scale=1.) # all_logits.append(torch.stack((output, ood_logits), dim=1)) if self.cfg.TRAINER.OOD_INFER_INTEGRATE: id_score = F.softmax(torch.stack((output, ood_logits), dim=1), dim=1)[:, 0, :] output *= id_score.clamp(min=min_thresh) else: ood_logits = None if self.cfg.TRAINER.OOD_INFER_OPTION == 'resume_res': output = image_features @ merged_text_features.t() if hasattr(self.dm.dataset, 'valid_classes'): output[:, ~self.dm.dataset.valid_classes] = -1. scores = get_msp_scores(output[:, self.dm.dataset.valid_classes]) else: scores, ood_flag = get_msp_scores(output, ood_logits, self.cfg.TRAINER.OOD_INFER, ret_near_ood=True) near_ood_flag.append(ood_flag) ood_score_list.append(scores.detach().cpu().numpy()) sc_labels_list.append(sc_labels.cpu().numpy()) ood_pred_list.append(output.argmax(dim=1).cpu().numpy()) ood_scores = np.concatenate(ood_score_list, axis=0) sc_labels = np.concatenate(sc_labels_list, axis=0) ood_preds = np.concatenate(ood_pred_list, axis=0) fake_ood_scores = ood_scores[sc_labels>=0] real_ood_scores = ood_scores[sc_labels<0] real_in_scores = np.concatenate([in_scores, fake_ood_scores], axis=0) if 'cifar' in self.dm.dataset.dataset_name: # compatible with SCOOD auroc, aupr, fpr95, thresh = get_measures(real_ood_scores, real_in_scores) else: # compatible with NPOS auroc, aupr, fpr95, thresh = get_measures(-real_in_scores, -real_ood_scores) print('auroc: %.4f, aupr: %.4f, fpr95: %.4f' % (auroc, aupr, fpr95)) save_lines.append('%10s auroc: %.4f, aupr: %.4f, fpr95: %.4f\n' % (ood_name, auroc, aupr, fpr95)) auroc_list.append(auroc) aupr_list.append(aupr) fpr95_list.append(fpr95) if self.cfg.TRAINER.OOD_PROMPT and len(near_ood_flag) and near_ood_flag[0] is not None: ood_tpr = torch.cat(near_ood_flag).sum().item() / len(ood_scores) print('NearOOD TPR: %.4f' % ood_tpr) ood_tpr_list.append(ood_tpr) if self.cfg.TRAINER.OOD_INFER_OPTION == 'save_res' and len(all_logits): torch.save(torch.cat(all_logits), f'{save_dir}/ood_{ood_name}_logits_all.pt') print('\nAverage: auroc: %.4f, aupr: %.4f, fpr95: %.4f' % (np.mean(auroc_list), np.mean(aupr_list), np.mean(fpr95_list))) save_lines.append('%10s auroc: %.4f, aupr: %.4f, fpr95: %.4f\n' % ('nAverage', np.mean(auroc_list), np.mean(aupr_list), np.mean(fpr95_list))) if self.cfg.TRAINER.OOD_PROMPT and len(ood_tpr_list) > 1: print('Average: OOD-TPR: %.4f' % np.mean(ood_tpr_list)) if model_directory != '': if 'ClassOOD' == ood_type: res_list = np.stack((auroc_list, aupr_list, fpr95_list), axis=1).reshape(-1,) * 100 np.savetxt(f'{model_directory}/{ood_type}_results.csv', res_list, fmt='%.2f', delimiter=',') save_path = f'{model_directory}/{ood_type}_results.txt' with open(save_path, 'w+') as f: f.writelines(save_lines) return list(results.values())[0], auroc, aupr, fpr95 @torch.no_grad() def cache_feat(self, split='train', is_ood=True): """A generic OOD testing pipeline.""" self.set_model_mode("eval") self.evaluator.reset() if split == 'train': data_loader = self.train_loader_x max_epoch = self.max_epoch else: data_loader = self.test_loader max_epoch = 1 if is_ood: from ood.datasets import LargeOODDataset from torch.utils.data import DataLoader data_root = osp.join(osp.abspath(osp.expanduser(self.cfg.DATASET.ROOT)), 'LargeOOD') for ood_name in ['inaturalist', 'sun', 'places', 'texture']: ood_set = LargeOODDataset(data_root, id_name=self.dm.dataset.dataset_name, ood_name=ood_name, transform=self.test_loader.dataset.transform) ood_loader = DataLoader(ood_set, batch_size=self.cfg.DATALOADER.TEST.BATCH_SIZE, shuffle=False, num_workers=self.test_loader.num_workers, drop_last=False, pin_memory=True) save_dir = f'{data_root}/clip_feat/{ood_name}' os.makedirs(save_dir, exist_ok=True) features, labels, paths = [], [], [] cnt = 0 for input, label in tqdm(ood_loader, desc='Caching image features'): input = input.to(self.device) label = label.to(self.device) image_features = self.model.image_encoder(input.type(self.model.dtype)).detach() features.append(image_features.cpu()) labels.append(label.cpu()) for i in range(len(input)): paths.append(ood_set.samples[cnt+i][0]) cnt += len(input) torch.save(torch.cat(features).half(), f'{save_dir}/test_image_features.pt') torch.save(torch.cat(labels).half(), f'{save_dir}/test_labels.pt') with open(f'{save_dir}/test_paths.txt', 'w+') as f: f.writelines([p + '\n' for p in paths]) else: save_dir = f'{self.dm.dataset.dataset_dir}/clip_feat' os.makedirs(save_dir, exist_ok=True) for self.epoch in range(self.start_epoch, max_epoch): features, labels, paths = [], [], [] for batch_idx, batch in enumerate(tqdm(data_loader, desc=f"Caching image features: {split} {self.epoch+1}/{max_epoch}: ")): input = batch["img"].to(self.device) label = batch["label"].to(self.device) image_features = self.model.image_encoder(input.type(self.model.dtype)).detach() features.append(image_features.cpu()) labels.append(label.cpu()) paths.extend(batch["impath"]) torch.save(torch.cat(features).half(), f'{save_dir}/ep{self.epoch}_{split}_image_features.pt') torch.save(torch.cat(labels).half(), f'{save_dir}/ep{self.epoch}_{split}_labels.pt') with open(f'{save_dir}/ep{self.epoch}_{split}_paths.txt', 'w+') as f: f.writelines([p + '\n' for p in paths])