main.py

# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))
# -*- coding: utf-8 -*-
# @Organization : Visual AI Lab, The University of Hong Kong
# @Author : Zeyu Chen && Jie Li
# @Function : Main codes for CodeBind

importlogging
importos
importdatetime
fromtypingimportAny
# log
try:
importwandb
exceptImportError:
wandb=None

try:
importmatplotlib.pyplotasplt
exceptImportError:
plt=None
logging.warning("Matplotlib not installed. This is not needed if you run this script as --headless")

importlightningasL
fromlightning.pytorchimportTrainer, seed_everything
fromlightning.pytorch.callbacksimportModelCheckpoint
fromlightning.pytorch.callbacksimportLearningRateMonitor
fromlightning.pytorchimportloggersaspl_loggers

importtorch
importtorch.nn.functionalasF
importtorch.optimasoptim
fromtorch.utils.dataimportDataLoader
fromtorchvision.utilsimportmake_grid

frommodelsimportcodebind_model
frommodelsimportloraasLoRA
frommodels.codebind_modelimportModalityType, load_module, save_module
fromapplication.emergent_zero_shot_classificationimportEmergentZeroShotClassifier, map_calculation
fromapplication.cross_modality_retrievalimportCrossModalityRetrieval

fromconfigimportload_cfg
fromdatasetsimportget_dataset
frommodels.VQimportbuild_vq
frommodels.decoderimportbuild_decoder


fromutilimportsave_image

importwarnings
warnings.filterwarnings('ignore')
logging.basicConfig(level=logging.INFO, force=True)

importpdb
fromicecreamimportic

# Logging settings
LOG_ON_STEP=True
LOG_ON_EPOCH=True
SAVE_CKP_LORA=True

MODALITIES= ["vision", "text", "audio", "thermal", "depth", "imu", "tactile", "eeg"]
fromdatasets.data_transform_rgbdimportinv_normalize

os.environ["CUBLAS_WORKSPACE_CONFIG"] =":4096:8"
torch.use_deterministic_algorithms(True, warn_only=True)


classCodeBind(L.LightningModule):
def__init__(self, args):
super().__init__()

# --- load args and set default valus
self.args=args
self.train_mode=args.get("train_mode", "lora") # train_mode in lora, fulltune, headtune
self.checkpoint_postfix=args.get("checkpoint_postfix", "_best")
self.checkpoint_dir=self.args.get("checkpoint_dir")
self.log_img_dir=self.args.get("log_img_dir")

self.modality_pair=args.modality_pair# 经过vq的模态
self.modality_eval=args.modality_eval# 待评估的目标模态
self.modality_train= [args.modality_train] ifisinstance(args.modality_train, str) elseargs.modality_train# 训练encoder的模态
ifself.modality_pair== ['vision', 'text'] orself.modality_pair== ['text', 'vision']:
self.modality_anchor, self.modality_nonanchor= ['vision'], 'text'
else:
self.modality_anchor=list(set(self.modality_pair) &set(['vision', 'text']))
self.modality_nonanchor=list(set(self.modality_pair) -set(self.modality_anchor))[0]
self.intra_anchor_align=args.get('intra_anchor_align', False) and'vision'inself.modality_anchorand'text'inself.modality_anchor
self.inter_anchor_align= [args.get('inter_anchor_align', None)] ifisinstance(args.get('inter_anchor_align', None), str) elseargs.get('inter_anchor_align', None)
self.inter_anchor_align=self.inter_anchor_alignifself.inter_anchor_alignisnotNoneelseself.modality_anchor

self.cfg_encoder=self.args.get('encoder')
self.cfg_decoder=self.args.get('decoder')

# set up datasets
self.train_dataset, self.test_dataset=get_dataset(args)

self.set_validation_metric(init=True)

self.save_hyperparameters()

self.init_network() 

# ---
self.ezs_classfier=None
self.retrieval=None
self.best_acc=0.
self.best_epoch=0
self.flag_update_text_model=TrueifModalityType.TEXTinself.modality_trainelseFalse
self.flag_update_retrieval_model=Trueifself.modality_evalinself.modality_trainelseFalse
print(f"self.flag_update_retrieval_model={self.flag_update_retrieval_model}")

ifself.args.get('load_checkpoint_dir') isnotNone:
self.load_flag=self.load_checkpoint(postfix=self.checkpoint_postfix, checkpoint_dir=self.args.get('load_checkpoint_dir'))
print(f"load_checkpoint all success = {self.load_flag} from {self.args.get('load_checkpoint_dir')}")

definit_network(self):
# --- Load full pretrained ImageBind model
self.model, self.model_postprocessors=codebind_model.imagebind_huge(pretrained=True, train_mode=self.train_mode, 
modality_train=self.modality_train, cfg_encoder=self.cfg_encoder, 
cfg_eeg=self.args.get('eeg_conf'))
logging.info(f"Enable {self.train_mode} for {self.modality_train}.")
ifself.train_mode=="lora":
lora_layer_idxs= {}
formodality_nameinself.args.lora_modality_names:
ifmodality_nameinMODALITIES:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs_{modality_name}', None)
ifnotlora_layer_idxs[modality_name]:
lora_layer_idxs[modality_name] =getattr(self.args, f'lora_layer_idxs', None)
else:
raiseValueError(f"Unknown modality name: {modality_name}")
logging.info("lora_layer_idxs:", lora_layer_idxs)
lora_module_mode=self.args.get("lora_module_mode", "attn_out")
logging.info(f"LoRA module mode: {lora_module_mode}")

# add lora weights in trunks
lora_model=LoRA.apply_lora_modality_trunks(self.model.modality_trunks, 
rank=self.args.get("lora_rank", 4),
layer_idxs=lora_layer_idxs,
modality_names=self.args.lora_modality_names,
lora_module_mode=lora_module_mode)

self.model.modality_trunks.update(lora_model)

# Decoder for reconstruction
self.use_decoder=self.cfg_decoder.get('use_decoder', True)
self.modality_reconstruction=self.args.get("modality_reconstruction", [])
self.modality_reconstruction= [self.modality_reconstruction] ifisinstance(self.modality_reconstruction, str) elseself.modality_reconstruction
self.modality_reconstruction=list(set(self.modality_pair).intersection(set(self.modality_reconstruction)))
ifself.use_decoder:
assertself.modality_reconstruction!= [], f"reconstruction modality should be specified if decoder is used"
self.decoder_trunks, self.decoder_heads=build_decoder(self.cfg_decoder, self.modality_reconstruction)
logging.info(f"Using reconstruction decoder for modality: {', '.join(self.modality_reconstruction)}")
else:
logging.info("Using reconstruction decoder for modality: None")

# multimodel vector quantiser
self.use_vq=self.args.get("vector_quantise", None) and (self.args.get("vector_quantise", None) isnotNone)
logging.info(f"Using vector quantise: {self.use_vq}")
ifself.use_vq:
self.cfg_vq=self.args.get('vector_quantise')
self.cfg_vq.update({'modality_pair': self.args.get('modality_pair')})
# pdb.set_trace()
self.modality_vq=build_vq(self.cfg_vq) 
else:
self.modality_vq=None

defset_validation_metric(self, init=False):
ifinit:
dataset_name=self.test_dataset[0].dataset_nameifisinstance(self.test_dataset, list) elseself.test_dataset.dataset_name
else:
dataset_name=self.trainer.val_dataloaders.dataset[0].dataset_nameifisinstance(self.trainer.val_dataloaders.dataset, list) elseself.trainer.val_dataloaders.dataset.dataset_name

ifdataset_name=='asa':
self.validation_metric='map'
elifdataset_namein ['msrvtt', 'audiocaps', 'clotho']:
self.validation_metric='recall'
else:
self.validation_metric='acc'
print(f"set_validation_metric '{self.validation_metric}' for {dataset_name}")

defconfigure_optimizers(self):
parameter_list= [
 {"params": self.model.modality_preprocessors.parameters()},
 {"params": self.model.modality_trunks.parameters()},
 {"params": self.model.modality_heads.parameters()},
 {"params": self.model.modality_postprocessors.parameters()},
 ]
lr_new=self.args.get('lr_new', self.args.lr)
# if self.args.get("sd_version", False):
# parameter_list.append({"params": self.model.mvq_adapter.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('add_new_output_head'):
parameter_list.append({"params": self.model.modality_heads_mvq.parameters(), "lr": lr_new})
ifself.cfg_encoder.get('use_postprocessors_outencoder'):
parameter_list.append({"params": self.model_postprocessors.parameters(), "lr": lr_new})
ifself.use_vq:
parameter_list.append({"params": self.modality_vq.parameters()})
ifself.use_decoder:
parameter_list.append({"params": self.decoder_trunks.parameters(), "lr": lr_new})
parameter_list.append({"params": self.decoder_heads.parameters(), "lr": lr_new})

optimizer=optim.AdamW(
parameter_list, 
lr=self.args.lr, 
weight_decay=self.args.weight_decay, 
betas=self.args.momentum_betas
 )

lr_scheduler=optim.lr_scheduler.CosineAnnealingLR(
optimizer, 
T_max=self.args.max_epochs, 
eta_min=self.args.lr/50
 )

return [optimizer], [lr_scheduler]

deftrain_dataloader(self):
ifisinstance(self.train_dataset, list) andlen(self.train_dataset) ==2:
train_loader=DataLoader(
self.train_dataset[self.current_epoch%2],
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers)
else:
train_loader=DataLoader(
self.train_dataset,
batch_size=self.args.batch_size,
shuffle=True,
drop_last=True,
pin_memory=False,
num_workers=self.args.num_workers,
 )
self.train_loader=train_loader
returntrain_loader

defval_dataloader(self):
ifisinstance(self.test_dataset, list):
val_loader=DataLoader(
self.test_dataset[self.current_epoch%len(self.test_dataset)],
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
else:
val_loader=DataLoader(
self.test_dataset,
batch_size=self.args.batch_size,
shuffle=False,
drop_last=False,
pin_memory=False,
num_workers=self.args.num_workers)
self.val_loader=val_loader
returnval_loader

defforward(self, batch, output_intermediate=False, modality_list=None):
ifmodality_listisNone:
modality_list=self.modality_pair
in_dict= {}
forclass_iinmodality_list:
ifbatch.get(class_i) isnotNone:
in_dict.update({class_i: batch.get(class_i)})
out_dict=self.model(in_dict, output_intermediate)

returnout_dict

defcalculate_loss(self, out_dict, batch, batch_idx, mode):
total_loss=0
ifmode=='train'orself.args.get('val_loss', False):
mode_vq='train'
if'text'inself.modality_pair:
vq_modality_list= [xforxinself.modality_pairifx!='text']
vq_modality_list.insert(0, 'text')
else:
vq_modality_list=self.modality_pair
else:
mode_vq='val'
vq_modality_list= [self.modality_eval]

feats_anchor= {modality_name: out_dict.get(modality_name) formodality_nameinself.modality_anchor}
feats_b= {self.modality_nonanchor: out_dict.get(self.modality_nonanchor)}

ifself.use_vq:
vq_dict= {}
ifself.cfg_vq.get('vq_all_token'): # mvq with cls_token and patch_tokens
formodality_nameinvq_modality_list:
ifmodality_name=='text':
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
else:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name+'_all_tokens'), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
out_dict.update({modality_name+'_vq': vq_dict_unimodal.get('concat')})
# get common feature for alignment text 仅有cls_token, 无patch_token
ifmodality_nameinself.modality_anchor:
feats_a=vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')
feats_anchor.update({modality_name: feats_a})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')[:, 0, ...] ifmodality_name!='text'elsevq_dict_unimodal.get('common')

else: # mvq with only cls_token
formodality_nameinvq_modality_list:
vq_dict_unimodal=self.modality_vq(out_dict.get(modality_name), modality_name, mode=mode_vq)
vq_dict.update({modality_name: vq_dict_unimodal})
ifmodality_nameinself.modality_anchor:
feats_anchor.update({modality_name: vq_dict_unimodal.get('common')})
else:
feats_b[modality_name] =vq_dict_unimodal.get('common')

loss_vq=self.vq_loss(vq_dict, mode=mode)
total_loss+=loss_vq

elifself.cfg_encoder.get('add_new_output_head'):
common_dim=self.cfg_encoder.get('output_embed_dim').get('common')
feats_anchor= {key: value[:, :common_dim] forkey, valueinfeats_anchor.items()} 
feats_b= {key: value[:, :common_dim] forkey, valueinfeats_b.items()} 

dual_nll=self.info_nce_loss(feats_anchor, feats_b, mode=mode)
total_loss+=dual_nll

ifself.use_decoder:
loss_recons=self.rec_loss(out_dict, batch, batch_idx, mode=mode)
total_loss+=loss_recons


returntotal_loss


definfo_nce_loss(self, feats_anchor, feats_b, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

# apply modality postprocessors if use it outside the encoder
ifself.cfg_encoder.get('use_postprocessors_outencoder', False):
assertself.model_postprocessorsisnotNone
formodality_name, featsinfeats_anchor.items():
# feats = F.normalize(feats, dim=-1)
feats_anchor.update({modality_name: self.model_postprocessors[modality_name](feats)})
formodality_name, featsinfeats_b.items():
# feats = F.normalize(feats, dim=-1)
feats_b.update({modality_name: self.model_postprocessors[modality_name](feats)})

ifself.args.self_contrastand'vision'inself.inter_anchor_align: # only modality vision has self_contrast
# TODO # both modality 'a' and 'b' have self_contrast
feats_tensors= [feats_anchor.get('vision').chunk(2)] + \
 [(feats_a.chunk(2)[0], feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [1] + [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= ["self"] + [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]
else:
feats_tensors= [(feats_a, feats_b[self.modality_nonanchor]) formodalname, feats_ainfeats_anchor.items() ifmodalnameinself.inter_anchor_align]
temperatures= [self.args.temperaturefor_inrange(len(self.inter_anchor_align))]
contrast= [f"cross_{modality_name}"formodality_nameinself.inter_anchor_align]

# intra anchor alignment when vision and text are all anchor modalities
ifself.intra_anchor_align:
feats_tensors+= [(feats_anchor.get('vision'), feats_anchor.get('text'))]
temperatures+= [self.args.temperature]
contrast+= ["intra"]

# Accumulate self-contrastive loss for image and its augmentation, and modailty with image
total_nll=0
forfeats_idx, (_feats_a_tensor, _feats_b_tensor) inenumerate(feats_tensors):
cos_sim=F.cosine_similarity(_feats_a_tensor[:, None, :], _feats_b_tensor[None, :, :], dim=-1)
cos_sim=cos_sim/temperatures[feats_idx]
nll_col=-torch.log(torch.diag(torch.softmax(cos_sim, dim=0)))
nll_row=-torch.log(torch.diag(torch.softmax(cos_sim, dim=1)))
nll= (nll_col+nll_row).mean()

total_nll+=nll

# Logging loss 
flag_acc=Falseifmode=='train'elseTrue# 训练时，acc_top1基本都为1，因此关闭以减少训练时间
ifflag_accand"cross"incontrast[feats_idx]:
# Get ranking position of positive example
pos_mask=torch.eye(cos_sim.shape[0], dtype=torch.bool, device=cos_sim.device)
comb_sim=torch.cat(
 [cos_sim[pos_mask][:, None], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first column) is positive example
dim=1,
 )
sim_argsort_1=comb_sim.argsort(dim=1, descending=True).argmin(dim=1) # find the similarity rank of positive example within the row of similarity matrix
comb_sim=torch.cat(
 [cos_sim[pos_mask][None, :], cos_sim.masked_fill(pos_mask, -9e15)], # First position (first row) is positive example
dim=0,
 )
sim_argsort_2=comb_sim.argsort(dim=0, descending=True).argmin(dim=0) # find the similarity rank of positive example within the column of similarity matrix

# Logging ranking metrics
self.log(mode+contrast[feats_idx] +"_acc_top1", ((sim_argsort_1==0).float().mean() + (sim_argsort_2==0).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
self.log(mode+contrast[feats_idx] +"_acc_top5", ((sim_argsort_1<5).float().mean() + (sim_argsort_2<5).float().mean()) /2,
prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True)
total_nll/=len(contrast)
self.log(mode+"_loss_cross", total_nll, prog_bar=True,
on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH,
sync_dist=True, # sync logging across all GPU workers,
# It is recommended to use 'sync_dist=True' when logging on epoch level in distributed setting to accumulate the metric across devices.
 )
returntotal_nll

defcmcm_loss(self, vq_dict):

loss_cmcm=0
modality_pair_list= [(modal_anchor, self.modality_nonanchor) formodal_anchorinself.modality_anchor]
ifself.intra_anchor_align:
modality_pair_list+= [('vision', 'text')]
for (modal_a, modal_b) inmodality_pair_list:
dist_a=vq_dict.get(modal_a).get('res_common').get('dist')
dist_b=vq_dict.get(modal_b).get('res_common').get('dist') # shape: vq_all_token: [bs, token_num, orig_feature_dim/common_embed_dim, common_embed_num]
# only cls token: [bs, orig_feature_dim/common_embed_dim, common_embed_num]

# 受限于显存，仅对cls_token 计算cmcm loss, 若注释，则对所有token的平均计算cmcm loss
dist_a=dist_a[:, 0, ...] ifdist_a.dim() ==4elsedist_a
dist_b=dist_b[:, 0, ...] ifdist_b.dim() ==4elsedist_b
# dist.dim() == 4: [bs, token_num, feature_num, embed_num] -> [token_num, bs, feature_num, embed_num]
# dist.dim() == 3: [bs, feature_num, embed_num] -> [bs, feature_num, embed_num]
dist_a=dist_a.permute(1,0,2,3) ifdist_a.dim() ==4elsedist_a
dist_b=dist_b.permute(1,0,2,3) ifdist_b.dim() ==4elsedist_b

# cmcm loss to align codebook across modalities
# probability distribution of each code measured by the distance of the code to every feature
iftorch.any(dist_a<0) ortorch.any(dist_b<0):
code_prob_a=F.softmax(dist_a, dim=-1) # [bs, feature_num, embed_num] or [token_num, bs, feature_num, embed_num]
code_prob_b=F.softmax(dist_b, dim=-1) 
else:
code_prob_a=F.softmax(torch.sqrt(dist_a), dim=1) 
code_prob_b=F.softmax(torch.sqrt(dist_b), dim=1)

# average for cls+patch tokens code prob
code_prob_a=code_prob_a.mean(0) ifcode_prob_a.dim() ==4elsecode_prob_a
code_prob_b=code_prob_b.mean(0) ifcode_prob_b.dim() ==4elsecode_prob_b

# [bs, feature_num, embed_num] * 2 -> [feature_num, bs, embed_num] * 2 -> [feature_num, bs, bs]
code_prob_a=code_prob_a.permute(1,0,2)
code_prob_b=code_prob_b.permute(1,0,2)
code_similarity=torch.einsum('imd,ind->imn', code_prob_a, torch.log(code_prob_b+1e-10)) +torch.einsum('imd,ind->imn', torch.log(code_prob_a+1e-10), code_prob_b)
code_similarity=code_similarity+torch.max(-code_similarity)
target_idx=torch.range(0, code_similarity.size(1)-1, dtype=torch.long, device=code_similarity.device).repeat(code_similarity.size(0))
loss_cmcm+=F.cross_entropy(code_similarity.view(-1, code_similarity.size(-1)), target_idx)

loss_cmcm/=len(modality_pair_list)
returnloss_cmcm


defmodal_decomp_loss(self, vq_dict):
# pdb.set_trace()
loss_modal_decomp=0
cnt=0
# vq_specific = []
formodality_name, vq_dict_unimodalinvq_dict.items():
common_unimodal=vq_dict_unimodal.get('common')
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
common_unimodal=F.normalize(common_unimodal, dim=-1)
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
loss_modal_decomp+=torch.norm(torch.einsum('bnd, bmd->bnm', common_unimodal, specific_unimodal), dim=(1,2), p='fro').mean()
cnt+=1
loss_modal_decomp/=cnt

returnloss_modal_decomp

defuniform_loss(self, vq_dict):
loss_uniform=0
cnt=0
formodality_name, vq_dict_unimodalinvq_dict.items():
specific_unimodal=vq_dict_unimodal.get('specific')
ifmodality_name!='text':
# only apply to cls token
specific_unimodal=specific_unimodal[:,0,...]
# perform normalization
ifnotself.cfg_vq.get('norm_code', False):
specific_unimodal=F.normalize(specific_unimodal, dim=-1)
# l2 distances among specific embeds
loss_uniform+=torch.pdist(specific_unimodal,p=2).pow(2).mul(-1).exp().mean().log()
cnt+=1

loss_uniform/=cnt

returnloss_uniform



defvq_loss(self, vq_dict, mode="train"):

ifmode=="val"andself.args.get('val_loss', False) isFalse:
return0

total_loss_vq=0
loss_vq, loss_vq_contra, loss_vq_uni=0, 0, 0
forvq_dict_unimodalinvq_dict.values():
loss_vq+=vq_dict_unimodal.get('loss_codebook')
loss_vq_contra+=vq_dict_unimodal.get('loss_contra')
loss_vq_uni+=vq_dict_unimodal.get('loss_uni')
total_loss_vq+=self.args.get('loss_vq_weight', 1.0) *loss_vq
total_loss_vq+=self.args.get('loss_vq_reg_weight', 1.0) * (loss_vq_contra+loss_vq_uni)

self.log(mode+"_loss_vq", loss_vq, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
ifself.args.get('loss_vq_reg_weight', 1.0) !=0.0:
self.log(mode+"_loss_vq_contra", loss_vq_contra, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
self.log(mode+"_loss_vq_uni", loss_vq_uni, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_cmcm_weight', 1.0) !=0.0:
loss_cmcm=self.cmcm_loss(vq_dict)
total_loss_vq+=self.args.get('loss_cmcm_weight', 1.0) *loss_cmcm
self.log(mode+"_loss_cmcm", loss_cmcm, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_modal_decomp_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_modal_decomp=self.modal_decomp_loss(vq_dict)
total_loss_vq+=self.args.get('loss_modal_decomp_weight', 1.0) *loss_modal_decomp
self.log(mode+"_loss_modal_decomp", loss_modal_decomp, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

ifself.args.get('loss_uniform_weight', 1.0) !=0.0andself.cfg_vq.get('codebook_num') >1:
loss_uniform=self.uniform_loss(vq_dict)
total_loss_vq+=self.args.get('loss_uniform_weight', 1.0) *loss_uniform
self.log(mode+"_loss_uniform", loss_uniform, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)
returntotal_loss_vq

defrec_loss(self, encoder_output, batch, batch_idx, mode='train'):
ifmode=="val"and (notself.args.get('val_loss', False)) andbatch_idx!=0:
return0

loss_piexl_recons=0
forclass_decoderinself.modality_reconstruction:
# class_encoder = 'vision' if class_decoder in ('image', 'video') else class_decoder 
ori_data=batch.get(class_decoder)
# encoder_trunk_input = encoder_output.get(class_encoder+"_trunk_input")
encoder_trunk_output=encoder_output.get(class_decoder+"_trunk_output")
encoder_pos_embed=encoder_output.get(class_decoder+"_pos_embed")
ifencoder_output.get(class_decoder+"_vq") isnotNone:
encoder_embedding=encoder_output.get(class_decoder+"_vq")
else:
encoder_embedding=encoder_output.get(class_decoder+"_all_tokens")

decoder_input=encoder_embeddingifself.cfg_decoder.get('with_cls_token') elseencoder_embedding[:, 1:, ...]

rec_token=self.decoder_trunks[class_decoder](
decoder_input,
orig_input_shape=None, 
input_pos_embed=encoder_pos_embed, 
input_trunk_embed=encoder_trunk_output,
use_checkpoint=False,
 )
# pdb.set_trace()
# ori_token = encoder_trunk_input[:, 1:, ...].transpose(1, 2)
rec_token=rec_token.transpose(1, 2) # channel first: [bs, token_num, token_dim] --> [bs, token_dim, token_num]. [bs, num=256, dim=512]
rec_tensor, ori_tensor=self.decoder_heads[class_decoder](rec_token, ori_data) # rec_img_h = ori_img_h / 2

# -- pixel-level reconstruction loss
ifmode=='train'or (mode=='val'andself.args.get('val_loss', False)):
class_decoder='vision'ifclass_decoderin ['image', 'video'] elseclass_decoder
loss_piexl_recons_x=F.mse_loss(rec_tensor, ori_tensor) 
self.log(mode+"_loss_piexl_recons_"+class_decoder, loss_piexl_recons_x, prog_bar=True, on_step=LOG_ON_STEP, on_epoch=LOG_ON_EPOCH, sync_dist=True)

loss_piexl_recons+=loss_piexl_recons_x


# visualize reconstruction results for some data in the first batch in both train and validation mode
ifbatch_idx==0: # if mode == "val" and batch_idx == 0:
ifori_tensor.ndim==5: # vision: only visualise frame 0 in video data
bs_ori=inv_normalize(class_decoder, ori_tensor[:, :, 0, ...]) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor[:, :, 0, ...].clone().detach())
else:
bs_ori=inv_normalize(class_decoder, ori_tensor) # value_range=(0, 1)
bs_rec=inv_normalize(class_decoder, rec_tensor.clone().detach())
max_images=min(bs_ori.size(0), 8)
img_pair=torch.cat((bs_ori[0:max_images,...], bs_rec[0:max_images,...]), dim=0)
# log image to disk
input_grid=make_grid(img_pair, nrow=max_images, value_range=(0, 1), normalize=True) # range -> value_range
_save_img_dir=os.path.join(self.log_img_dir, mode)
os.makedirs(_save_img_dir, exist_ok=True)
log_img_path=os.path.join(_save_img_dir, f'ori_rec_{class_decoder}_epoch{self.current_epoch}.jpg')
save_image(255*input_grid.cpu().numpy(), image_path=log_img_path)

returnloss_piexl_recons

defon_train_epoch_start(self) ->None:
returnsuper().on_train_epoch_start()

deftraining_step(self, batch, batch_idx):
output_intermediate=True# if self.use_decoder else False
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="train")

returntotal_loss

defon_train_epoch_end(self):
pass


defvalidation_step(self, batch, batch_idx):

ifself.validation_metric=='map':
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
audio_outputs, labels=self.ezs_classfier.get_batch_classification_map(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs['audio_outputs'].append(audio_outputs)
self.validation_step_outputs['labels'].append(labels)

self.number_of_examples+=data_b.size(0)

elifself.validation_metric=='recall':
gt_name=batch.get('retrieval_id')
bs_recall=self.retrieval.get_batch_recall(batch.get('text'), 'text', gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_recall)

self.number_of_examples+=batch.get('text').size(0)
else:
gt_name=batch.get('label')
data_b=batch.get(self.modality_eval)
class_b=self.modality_eval
# pdb.set_trace()
bs_acc=self.ezs_classfier.get_batch_classification_acc(data_b, class_b, gt_name, vector_quantiser=self.modality_vq)
self.validation_step_outputs.append(bs_acc)

self.number_of_examples+=data_b.size(0)

ifself.args.get('val_loss', False) or (self.use_decoderandbatch_idx==0):
# pdb.set_trace()
output_intermediate=True
res_dict=self.forward(batch, output_intermediate)
total_loss=self.calculate_loss(res_dict, batch, batch_idx, mode="val")

defon_validation_epoch_start(self) ->None:
self.set_validation_metric()
ifself.validation_metric=='recall':
ifself.flag_update_retrieval_modelorself.retrievalisNone:
print(f"Prepare Cross Madality Retrieval for {self.modality_eval} ...")
self.retrieval=CrossModalityRetrieval(classifier_model=self, retrieval_modality_type=self.modality_eval)
else:
ifself.use_vqorself.flag_update_text_modelorself.ezs_classfierisNone:
print(f"Prepare Emergent ZeroShot Classifier for {self.modality_eval} ... ")
self.ezs_classfier=EmergentZeroShotClassifier(classifier_model=self, modality_type=self.modality_eval, text_template=self.args.get('text_template'))

ifself.validation_metric=='map':
self.validation_step_outputs= {'audio_outputs': [], 'labels': []} 
else:
self.validation_step_outputs= []
self.number_of_examples=0
returnsuper().on_validation_epoch_start()

defon_validation_epoch_end(self): 
# -- for data in each rank:

# -- for data in all ranks:
all_val_out=self.all_gather(self.validation_step_outputs)
all_number_of_examples=self.all_gather(self.number_of_examples)

ifself.trainer.is_global_zero:
# merge output and process
ifself.validation_metric=='map':
# use map as evaluation metric
all_num=sum(all_number_of_examples)
audio_outputs= []
forbatch_val_outinall_val_out['audio_outputs']:
ifbatch_val_out.dim() ==3:
audio_outputs.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
audio_outputs.append(batch_val_out)
audio_outputs=torch.cat(audio_outputs).to('cpu').detach()
labels= []
forbatch_val_outinall_val_out['labels']: # label: GT index of all classes 
ifbatch_val_out.dim() ==3:
labels.append(batch_val_out.reshape(-1, batch_val_out.size(-1)))
else:
assertbatch_val_out.dim() ==2
labels.append(batch_val_out)
labels=torch.cat(labels).to('cpu').detach()
all_acc= [map_calculation(audio_outputs, labels)] # scale to list, to be compatible with other metrics

else:
# use accuracy or recall as evaluation metric
all_num=sum(all_number_of_examples) # list of int
all_val_out=sum(all_val_out) # list of list --> flattend list
all_acc=sum(all_val_out) /all_num*100.

ifself.best_acc<all_acc[0]:
self.best_acc=all_acc[0]
self.best_epoch=self.current_epoch
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_best')

logging.info(f"number_of_examples={all_num}, top 1 accuracy = {all_acc[0]:.2f}%, best_accuracy = {self.best_acc:.2f}% at epoch={self.best_epoch}")
 [print(f"Accuracy {i}{acc:.3f}%", end=" |") fori, accinenumerate(all_acc)]

self.log("val_acc", all_acc[0], on_epoch=LOG_ON_EPOCH, sync_dist=True, rank_zero_only=True)

self.validation_step_outputs.clear() # free memory
self.trainer.strategy.barrier() # to let other cards to wait
ifSAVE_CKP_LORA:
self.save_checkpoint(postfix='_last')

defload_checkpoint(self, postfix='_best', checkpoint_dir=None):
ifcheckpoint_dirisNone:
checkpoint_dir=self.checkpoint_dir

load_flag= []
# Load existing preprocessors & heads & postprocessors
load_flag.append(load_module(self.model.modality_preprocessors, modality_name=self.modality_pair,
module_name="preprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_heads, modality_name=self.modality_pair,
module_name="heads", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.model.modality_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.model_postprocessorsisnotNone:
load_flag.append(load_module(self.model_postprocessors, modality_name=self.modality_pair,
module_name="postprocessors_new", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifhasattr(self.model, 'modality_heads_mvq'):
load_flag.append(load_module(self.model.modality_heads_mvq, modality_name=self.modality_pair,
module_name="heads_mvq", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifhasattr(self.model, 'mvq_adapter'):
load_flag.append(load_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=checkpoint_dir, postfix=postfix))
ifself.train_mode=="lora":
# use fixed trunks with loaded lora parameters
load_flag.append(LoRA.load_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=checkpoint_dir))
elifself.train_mode=="fulltune":
# Load existing trunks
load_flag.append(load_module(self.model.modality_trunks, modality_name=self.modality_pair,
module_name="trunks", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load codebook
ifself.use_vq:
load_flag.append(load_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=checkpoint_dir, postfix=postfix))

# load decoder
ifself.use_decoderandself.args.get('train', False): # test 时，不需要 decoder
load_flag.append(load_module(self.decoder_trunks, modality_name=self.modality_reconstruction,
module_name="decoder_trunk", checkpoint_dir=checkpoint_dir, postfix=postfix))
load_flag.append(load_module(self.decoder_heads, modality_name=self.modality_reconstruction,
module_name="decoder_head", checkpoint_dir=checkpoint_dir, postfix=postfix))

ifall(load_flag):
returnTrue
else:
print(f"load_flag status = {load_flag}")
returnFalse

defsave_checkpoint(self, postfix='_best'):

# Save preprocessors & heads & postprocessors
save_module(self.model.modality_preprocessors, modality_name=self.modality_train,
module_name="preprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_postprocessors, modality_name=self.modality_train,
module_name="postprocessors", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.model.modality_heads, modality_name=self.modality_train,
module_name="heads", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifself.model_postprocessorsisnotNone:
save_module(self.model_postprocessors, modality_name=self.modality_train,
module_name="postprocessors_new", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'modality_heads_mvq'):
save_module(self.model.modality_heads_mvq, modality_name=self.modality_train,
module_name="heads_mvq", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
ifhasattr(self.model, 'mvq_adapter'):
save_module(self.model.mvq_adapter, modality_name=['vision', 'text'],
module_name="mvq_adapter", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

ifself.train_mode=="lora":
LoRA.save_lora_modality_trunks(self.model.modality_trunks, checkpoint_dir=self.checkpoint_dir, postfix=postfix)
elifself.train_mode=="fulltune":
# Save preprocessors, trunks, heads, postprocessors in whole model fine-tuning situation
save_module(self.model.modality_trunks, modality_name=self.modality_train,
module_name="trunks", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save Vector Quantise
ifself.use_vq:
save_module(self.modality_vq, module_name="vector_quantise", checkpoint_dir=self.checkpoint_dir, postfix=postfix)

# save decoder
ifself.use_decoder:
save_module(self.decoder_trunks, modality_name=self.modality_reconstruction, 
module_name="decoder_trunk", checkpoint_dir=self.checkpoint_dir, postfix=postfix)
save_module(self.decoder_heads, modality_name=self.modality_reconstruction, 
module_name="decoder_head", checkpoint_dir=self.checkpoint_dir, postfix=postfix)


defset_logger(args):
# Create logger
logger=args.get("loggers", None)
iflogger=="wandb":
ifwandbisNone:
raiseImportError("wandb is not installed. Install it or choose another logger.")

save_dir=args.get('loggers_dir', f"./exp/{args.get('expname')}/log")
os.makedirs(save_dir, exist_ok=True)

wandb_project=args.get("wandb_project", "codebind")
wandb_entity=args.get("wandb_entity", "codebind")

wandb_logger=pl_loggers.WandbLogger(
save_dir=save_dir,
project=wandb_project,
entity=wandb_entity,
name=args.get('expname'))
wandb_logger.experiment.config.update(dict(args), allow_val_change=True)
logger=wandb_logger
else:
raiseNotImplementedError(f"Logger {logger} is not supported. Please choose from ['wandb'].")

print("Create logger:", args.get("loggers", None))
returnlogger


defmain(args):
# Set experiment properties
seed_everything(args.seed, workers=True)

model=CodeBind(args=args)
# load existing checkpoint
if (notargs.get('train', False)) and (notSAVE_CKP_LORA):
save_checkpoint_dir=os.path.join(args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"), "example_best.ckpt")
ifos.path.exists(save_checkpoint_dir):
model=CodeBind.load_from_checkpoint(save_checkpoint_dir)


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
train(args, model)
else:
logging.info("# # # # # ==> Running testing progress ...")
test(args, model)


deftrain(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
accelerator="cpu"
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name] 
devices.sort()
accelerator="gpu"
print("devices =", devices)

# whether to reload dataloaders between epochs based on number of validation datasets
ifisinstance(model.test_dataset, list):
reload_dataloaders_epoch_num=1
else:
reload_dataloaders_epoch_num=0

logger=set_logger(args)

checkpoint_callback=ModelCheckpoint(
dirpath=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
filename='example_best',
save_top_k=1,
verbose=True,
monitor='val_acc',
mode='max',
 )

logging.info("Prepare trainer")
trainer=Trainer(accelerator=accelerator,
devices=devices,
strategy='ddp_find_unused_parameters_true',
deterministic=True,
max_epochs=args.max_epochs, 
gradient_clip_val=args.gradient_clip_val, # callbacks=[lr_monitor]
check_val_every_n_epoch=1, # validation frequency
reload_dataloaders_every_n_epochs=reload_dataloaders_epoch_num,
limit_train_batches=args.get('limit_train_batches', None), # debug, run through only 10% of the training set each epoch
limit_val_batches=args.get('limit_val_batches', None),
logger=loggerifloggerelseNone, 
default_root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log"),
enable_checkpointing=FalseifSAVE_CKP_LORAelseTrue,
num_sanity_val_steps=0,
callbacks=[checkpoint_callback] ifnotSAVE_CKP_LORAelseNone
 )


ifargs.get('train', False):
logging.info("# # # # # ==> Running training progress ...") 
trainer.fit(model)
ifnotSAVE_CKP_LORA:
root_dir=args.get('lightning_log', f"./exp/{args.get('expname')}/lightning_log")
trainer.save_checkpoint(os.path.join(root_dir, "example_last.ckpt"))

else:
logging.info("# # # # # ==> Running testing progress ...")
trainer.validate(model)



deftest(args, model):
device_name=args.device
if"cpu"==device_name[0]:
devices=1
else:
devices= [int(device_name.split(":")[1])] ifisinstance(device_name, str) else [int(i.split(":")[1]) foriindevice_name]
devices.sort()
print("devices =", devices)
device=torch.device("cpu"if"cpu"==device_name[0] elsef"cuda:{devices[0]}")

# -----------------------------------
logging.info('Start evaluate Emergent Zero-Shot Classification')
model.val_dataloader()
model.to(device)
model.freeze()


ifmodel.validation_metric=='recall':
retrieval=CrossModalityRetrieval(classifier_model=model,
retrieval_modality_type=args.modality_eval,
retrieval_data=model.test_dataset.data_paths)
retrieval.get_retrieval_recall(model.val_loader)
else:
ezs_classfier=EmergentZeroShotClassifier(classnames=model.test_dataset.class_names, 
classifier_model=model,
dataset_name=args.datasets,
modality_type=args.modality_eval,
text_template=args.get('text_template'))

ifmodel.validation_metric=='acc':
ezs_classfier.get_classification_acc(model.val_loader)
else:
ezs_classfier.get_classification_map(model.val_loader)


if__name__=="__main__":
time_start=datetime.datetime.now()
cfg_all=load_cfg()
main(cfg_all)
time_end=datetime.datetime.now()
print('Running time: %s'%(time_end-time_start))