462 lines
19 KiB
Python
462 lines
19 KiB
Python
import copy
|
|
import os.path as osp
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
from torch.cuda.amp import GradScaler, autocast
|
|
import json
|
|
|
|
from dassl.engine import TRAINER_REGISTRY, TrainerX
|
|
from dassl.utils import load_pretrained_weights, load_checkpoint
|
|
from dassl.optim import build_optimizer, build_lr_scheduler
|
|
from clip import clip
|
|
from clip.simple_tokenizer import SimpleTokenizer as _Tokenizer
|
|
|
|
_tokenizer = _Tokenizer()
|
|
|
|
DESC_LLM = "gpt-4.1"
|
|
DESC_TOPK = 4
|
|
|
|
CUSTOM_TEMPLATES = {
|
|
"OxfordPets": "a photo of a {}, a type of pet.",
|
|
"OxfordFlowers": "a photo of a {}, a type of flower.",
|
|
"FGVCAircraft": "a photo of a {}, a type of aircraft.",
|
|
"DescribableTextures": "a photo of a {}, a type of texture.",
|
|
"EuroSAT": "a centered satellite photo of {}.",
|
|
"StanfordCars": "a photo of a {}.",
|
|
"Food101": "a photo of {}, a type of food.",
|
|
"SUN397": "a photo of a {}.",
|
|
"Caltech101": "a photo of a {}.",
|
|
"UCF101": "a photo of a person doing {}.",
|
|
"ImageNet": "a photo of a {}.",
|
|
"ImageNetSketch": "a photo of a {}.",
|
|
"ImageNetV2": "a photo of a {}.",
|
|
"ImageNetA": "a photo of a {}.",
|
|
"ImageNetR": "a photo of a {}.",
|
|
}
|
|
|
|
|
|
def load_clip_to_cpu(cfg, zero_shot_model=False):
|
|
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")
|
|
if not zero_shot_model:
|
|
design_details = {"trainer": 'IVLP',
|
|
"vision_depth": cfg.TRAINER.DZGCOOP.PROMPT_DEPTH_VISION,
|
|
"language_depth": cfg.TRAINER.DZGCOOP.PROMPT_DEPTH_TEXT,
|
|
"vision_ctx": cfg.TRAINER.DZGCOOP.N_CTX_VISION,
|
|
"language_ctx": cfg.TRAINER.DZGCOOP.N_CTX_TEXT}
|
|
model = clip.build_model(state_dict or model.state_dict(), design_details)
|
|
else:
|
|
# Return original CLIP model for generating frozen VL features
|
|
design_details = {"trainer": 'IVLP',
|
|
"vision_depth": 0,
|
|
"language_depth": 0, "vision_ctx": 0,
|
|
"language_ctx": 0}
|
|
model = clip.build_model(state_dict or model.state_dict(), design_details)
|
|
return model
|
|
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):
|
|
x = prompts + self.positional_embedding.type(self.dtype)
|
|
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 VLPromptLearner(nn.Module):
|
|
def __init__(self, cfg, classnames, clip_model):
|
|
super().__init__()
|
|
n_cls = len(classnames)
|
|
# Make sure Language depth >= 1
|
|
assert cfg.TRAINER.DZGCOOP.PROMPT_DEPTH_TEXT >= 1, "In Independent VL prompting, Language prompt depth should be >=1" \
|
|
"\nPlease use VPT trainer if you want to learn only vision " \
|
|
"branch"
|
|
n_ctx = cfg.TRAINER.DZGCOOP.N_CTX_TEXT
|
|
ctx_init = cfg.TRAINER.DZGCOOP.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})"
|
|
|
|
if ctx_init and n_ctx <= 4:
|
|
ctx_init = ctx_init.replace("_", " ")
|
|
prompt = clip.tokenize(ctx_init)
|
|
with torch.no_grad():
|
|
embedding = clip_model.token_embedding(prompt).type(dtype)
|
|
ctx_vectors_strong = embedding[0, 1: 1 + n_ctx, :]
|
|
prompt_prefix_strong = ctx_init
|
|
else:
|
|
ctx_vectors_strong = torch.empty(n_ctx, ctx_dim, dtype=dtype)
|
|
nn.init.normal_(ctx_vectors_strong, std=0.02)
|
|
prompt_prefix_strong = " ".join(["X"] * n_ctx)
|
|
|
|
ctx_vectors_weak = torch.empty(n_ctx, ctx_dim, dtype=dtype)
|
|
nn.init.normal_(ctx_vectors_weak, std=0.02)
|
|
prompt_prefix_weak = " ".join(["X"] * n_ctx)
|
|
|
|
print(f"Independent V-L design with Dual Prompt Branches")
|
|
print(f'Strong branch initial text context: "{prompt_prefix_strong}"')
|
|
print(f'Weak branch initial text context: "{prompt_prefix_weak}"')
|
|
print(f"Number of context words (tokens) for Language prompting: {n_ctx}")
|
|
print(f"Number of context words (tokens) for Vision prompting: {cfg.TRAINER.DZGCOOP.N_CTX_VISION}")
|
|
self.ctx_strong = nn.Parameter(ctx_vectors_strong)
|
|
self.ctx_weak = nn.Parameter(ctx_vectors_weak)
|
|
|
|
classnames = [name.replace("_", " ") for name in classnames]
|
|
name_lens = [len(_tokenizer.encode(name)) for name in classnames]
|
|
prompts = [prompt_prefix_strong + " " + name + "." for name in classnames]
|
|
|
|
tokenized_prompts = torch.cat([clip.tokenize(p) for p in prompts]) # (n_cls, n_tkn)
|
|
# Also create frozen CLIP
|
|
clip_model_temp = load_clip_to_cpu(cfg, True).float().cuda()
|
|
clip_model_temp_image = load_clip_to_cpu(cfg, True)
|
|
with torch.no_grad():
|
|
embedding = clip_model.token_embedding(tokenized_prompts).type(dtype)
|
|
self.ZS_image_encoder = clip_model_temp_image.visual
|
|
# Now pre-compute the frozen VL embeddings from LLM descriptions
|
|
semantic_guidance_features = []
|
|
desc_file = f"./desc/{DESC_LLM}/descriptions_top{DESC_TOPK}/{cfg.DATASET.NAME}.json"
|
|
with open(desc_file, "r") as f:
|
|
all_desc = json.load(f)
|
|
|
|
template = CUSTOM_TEMPLATES[cfg.DATASET.NAME]
|
|
for cls in classnames:
|
|
cls_descs = [template.format(cls)[:-1] + f", {desc}" for desc in all_desc[cls]]
|
|
cls_token = torch.cat([clip.tokenize(cls_desc) for cls_desc in cls_descs]).cuda()
|
|
with torch.no_grad():
|
|
cls_feature = clip_model_temp.encode_text(cls_token)
|
|
cls_feature = cls_feature / cls_feature.norm(dim=-1, keepdim=True)
|
|
cls_feature = torch.mean(cls_feature, dim=0)
|
|
semantic_guidance_features.append(cls_feature)
|
|
|
|
self.semantic_embeddings = torch.stack(semantic_guidance_features)
|
|
print(f"Using LLM descriptions from: {desc_file}")
|
|
# 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
|
|
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
|
|
|
|
def construct_prompts(self, ctx, prefix, suffix, label=None):
|
|
# dim0 is either batch_size (during training) or n_cls (during testing)
|
|
# ctx: context tokens, with shape of (dim0, n_ctx, ctx_dim)
|
|
# prefix: the sos token, with shape of (n_cls, 1, ctx_dim)
|
|
# suffix: remaining tokens, with shape of (n_cls, *, ctx_dim)
|
|
|
|
if label is not None:
|
|
prefix = prefix[label]
|
|
suffix = suffix[label]
|
|
|
|
prompts = torch.cat(
|
|
[
|
|
prefix, # (dim0, 1, dim)
|
|
ctx, # (dim0, n_ctx, dim)
|
|
suffix, # (dim0, *, dim)
|
|
],
|
|
dim=1,
|
|
)
|
|
|
|
return prompts
|
|
|
|
def forward(self):
|
|
ctx_strong = self.ctx_strong
|
|
ctx_weak = self.ctx_weak
|
|
|
|
if ctx_strong.dim() == 2:
|
|
ctx_strong = ctx_strong.unsqueeze(0).expand(self.n_cls, -1, -1)
|
|
ctx_weak = ctx_weak.unsqueeze(0).expand(self.n_cls, -1, -1)
|
|
|
|
prefix = self.token_prefix
|
|
suffix = self.token_suffix
|
|
prompts_strong = self.construct_prompts(ctx_strong, prefix, suffix)
|
|
prompts_weak = self.construct_prompts(ctx_weak, prefix, suffix)
|
|
|
|
return prompts_strong, prompts_weak
|
|
|
|
|
|
class CustomCLIP(nn.Module):
|
|
def __init__(self, cfg, classnames, clip_model):
|
|
super().__init__()
|
|
self.prompt_learner = VLPromptLearner(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
|
|
self.dtype = clip_model.dtype
|
|
self.total_epochs = cfg.OPTIM.MAX_EPOCH
|
|
self.n_cls = len(classnames)
|
|
|
|
def forward(self, image, label=None):
|
|
tokenized_prompts = self.tokenized_prompts
|
|
logit_scale = self.logit_scale.exp()
|
|
|
|
prompts_strong, prompts_weak = self.prompt_learner()
|
|
|
|
with torch.no_grad():
|
|
zero_shot_features = self.prompt_learner.ZS_image_encoder(image.type(self.dtype))
|
|
zero_shot_features = zero_shot_features / zero_shot_features.norm(dim=-1, keepdim=True)
|
|
|
|
image_features = self.image_encoder(image.type(self.dtype))
|
|
image_features = image_features / image_features.norm(dim=-1, keepdim=True)
|
|
|
|
text_features_strong = self.text_encoder(prompts_strong, tokenized_prompts)
|
|
text_features_strong = text_features_strong / text_features_strong.norm(dim=-1, keepdim=True)
|
|
|
|
text_features_weak = self.text_encoder(prompts_weak, tokenized_prompts)
|
|
text_features_weak = text_features_weak / text_features_weak.norm(dim=-1, keepdim=True)
|
|
|
|
semantic_embeddings = self.prompt_learner.semantic_embeddings
|
|
semantic_embeddings = semantic_embeddings / semantic_embeddings.norm(dim=-1, keepdim=True)
|
|
|
|
zero_shot_logits = logit_scale * zero_shot_features.cuda() @ semantic_embeddings.half().cuda().t()
|
|
|
|
logits_strong = logit_scale * image_features @ text_features_strong.t()
|
|
logits_weak = logit_scale * image_features @ text_features_weak.t()
|
|
|
|
zs_probs = F.softmax(zero_shot_logits, dim=1)
|
|
confidence = zs_probs.max(dim=1).values
|
|
|
|
alpha = confidence.unsqueeze(1)
|
|
|
|
logits_final = alpha * logits_strong + (1 - alpha) * logits_weak
|
|
|
|
if self.prompt_learner.training:
|
|
loss_ce = F.cross_entropy(logits_final, label)
|
|
return loss_ce, text_features_strong, text_features_weak, semantic_embeddings, zero_shot_features, image_features, zero_shot_logits, logits_strong, logits_weak, logits_final
|
|
else:
|
|
return logits_final
|
|
|
|
|
|
@TRAINER_REGISTRY.register()
|
|
class DZGCoOp(TrainerX):
|
|
def check_cfg(self, cfg):
|
|
assert cfg.TRAINER.DZGCOOP.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.DZGCOOP.PREC == "fp32" or cfg.TRAINER.DZGCOOP.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")
|
|
name_to_update = "prompt_learner"
|
|
|
|
for name, param in self.model.named_parameters():
|
|
if name_to_update not in name:
|
|
# Make sure that VPT prompts are updated
|
|
if "VPT" in name:
|
|
param.requires_grad_(True)
|
|
else:
|
|
param.requires_grad_(False)
|
|
else:
|
|
if "ZS_image_encoder" in name:
|
|
param.requires_grad_(False)
|
|
|
|
# Double check
|
|
enabled = set()
|
|
for name, param in self.model.named_parameters():
|
|
if param.requires_grad:
|
|
enabled.add(name)
|
|
print(f"Parameters to be updated: {enabled}")
|
|
print(f"Parameters count: {len(enabled)}")
|
|
if cfg.MODEL.INIT_WEIGHTS:
|
|
load_pretrained_weights(self.model, cfg.MODEL.INIT_WEIGHTS)
|
|
|
|
self.model.to(self.device)
|
|
# NOTE: only give prompt_learner to the optimizer
|
|
self.optim = build_optimizer(self.model, cfg.OPTIM)
|
|
self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
|
|
self.register_model("VLPromptLearner", self.model, self.optim, self.sched)
|
|
# Cosine scheduler
|
|
self.total_epochs = cfg.OPTIM.MAX_EPOCH
|
|
self.step_counter = 1
|
|
N = cfg.OPTIM.MAX_EPOCH
|
|
mean = cfg.TRAINER.DZGCOOP.EWA_MEAN
|
|
stdev = cfg.TRAINER.DZGCOOP.EWA_STD
|
|
normal = self.get_normal(mean, stdev)
|
|
self.normal_weights = np.array([normal(a) for a in range(1, N + 1)])
|
|
self.normal_weights = self.normal_weights / sum(self.normal_weights)
|
|
self.scaler = GradScaler() if cfg.TRAINER.DZGCOOP.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)
|
|
# Keep model with EWA
|
|
self.previous_model_ewa = None
|
|
|
|
def forward_backward(self, batch):
|
|
image, label = self.parse_batch_train(batch)
|
|
|
|
model = self.model
|
|
optim = self.optim
|
|
scaler = self.scaler
|
|
|
|
prec = self.cfg.TRAINER.DZGCOOP.PREC
|
|
if prec == "amp":
|
|
with autocast():
|
|
loss = model(image, label)
|
|
optim.zero_grad()
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optim)
|
|
scaler.update()
|
|
else:
|
|
loss_ce, text_features_strong, text_features_weak, semantic_embeddings, zs_image_embedd, image_ft, \
|
|
zero_shot_logits, logits_strong, logits_weak, logits_final = model(image, label)
|
|
|
|
lambda1 = self.cfg.TRAINER.DZGCOOP.IMAGE_LOSS_WEIGHT
|
|
lambda2 = self.cfg.TRAINER.DZGCOOP.TEXT_LOSS_WEIGHT_STRONG
|
|
lambda3 = self.cfg.TRAINER.DZGCOOP.TEXT_LOSS_WEIGHT_WEAK
|
|
|
|
L_zvg = F.l1_loss(image_ft, zs_image_embedd.cuda(), reduction='mean') * lambda1
|
|
L_sg_strong = F.l1_loss(text_features_strong, semantic_embeddings.cuda(), reduction='mean') * lambda2
|
|
L_sg_weak = F.l1_loss(text_features_weak, semantic_embeddings.cuda(), reduction='mean') * lambda3
|
|
|
|
L_zpg = F.kl_div(
|
|
F.log_softmax(logits_final / 1, dim=1),
|
|
F.log_softmax(zero_shot_logits / 1, dim=1),
|
|
reduction='sum',
|
|
log_target=True
|
|
) * (1 * 1) / logits_final.numel()
|
|
|
|
L_zg = (L_zpg + L_sg_strong + L_sg_weak + L_zvg)
|
|
loss = (loss_ce + L_zg)
|
|
optim.zero_grad()
|
|
loss.backward()
|
|
optim.step()
|
|
|
|
loss_summary = {"loss": loss.item()}
|
|
|
|
if (self.batch_idx + 1) == self.num_batches:
|
|
self.update_lr()
|
|
# Means one epoch is completed, perform EWA
|
|
self.step_counter = self.step_counter + 1
|
|
current_epoch_weight = self.normal_weights[self.step_counter - 2]
|
|
current_model_weights = copy.deepcopy(model.state_dict())
|
|
for key in current_model_weights:
|
|
current_model_weights[key] = current_model_weights[key].cpu()
|
|
weighted_state_dict = self.state_dict_weighting(current_model_weights, current_epoch_weight)
|
|
if self.previous_model_ewa is None:
|
|
self.previous_model_ewa = weighted_state_dict
|
|
else:
|
|
self.previous_model_ewa = self.state_dict_add(weighted_state_dict, self.previous_model_ewa)
|
|
|
|
if self.step_counter == self.model.total_epochs + 1:
|
|
print("Using EWA model for final inference...")
|
|
model.load_state_dict(self.previous_model_ewa)
|
|
self.model.load_state_dict(self.previous_model_ewa)
|
|
return loss_summary
|
|
|
|
def state_dict_weighting(self, main_dict, weightage, prompt_only=False):
|
|
# Average all parameters
|
|
updated_dict = copy.deepcopy(main_dict)
|
|
if not prompt_only:
|
|
for key in main_dict:
|
|
updated_dict[key] = main_dict[key].cpu() * weightage
|
|
return updated_dict
|
|
else:
|
|
return main_dict.cpu() * weightage
|
|
|
|
def state_dict_add(self, dict1, dict2, prompt_only=False):
|
|
# Average all parameters
|
|
if not prompt_only:
|
|
modified_dict = dict2
|
|
for key in dict1:
|
|
modified_dict[key] = modified_dict[key].cpu() + dict1[key].cpu()
|
|
return modified_dict
|
|
else:
|
|
return dict1.cpu() + dict2.cpu()
|
|
|
|
def get_normal(self, mu, sigma):
|
|
normal = lambda x: (1 / (sigma * np.sqrt(2 * np.pi))) * np.exp(-0.5 * ((x - mu) / sigma) ** 2)
|
|
return normal
|
|
|
|
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
|
|
if "prompt_learner.token_prefix" in state_dict:
|
|
del state_dict["prompt_learner.token_prefix"]
|
|
|
|
if "prompt_learner.token_suffix" in state_dict:
|
|
del state_dict["prompt_learner.token_suffix"]
|
|
|
|
# Handle backward compatibility: if old checkpoint has ctx, initialize both ctx_strong and ctx_weak
|
|
if "prompt_learner.ctx" in state_dict:
|
|
ctx = state_dict.pop("prompt_learner.ctx")
|
|
state_dict["prompt_learner.ctx_strong"] = ctx.clone()
|
|
state_dict["prompt_learner.ctx_weak"] = ctx.clone()
|
|
|
|
print("Loading weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch))
|
|
# set strict=False
|
|
self._models[name].load_state_dict(state_dict, strict=False)
|