import copy import datetime import logging import os import time from os.path import join import pandas as pd import torch import torch.backends.cudnn as cudnn import torch.distributed as dist import wandb from dataset import MetaLoader #from models.vindlu import VindLU #from models.vindlu_vit import VindLU_VIT #from models.vindlu_videoclip import VindLU_VideoCLIP #from models.vindlu_blip_qformer import VindLU_BLIP_QFormer from models.viclip import ViCLIP from tasks.pretrain import setup_dataloaders from tasks.retrieval_utils import evaluation_wrapper from tasks.shared_utils import setup_model from utils.basic_utils import MetricLogger, SmoothedValue, setup_seed from utils.config import Config from utils.config_utils import setup_main from utils.distributed import get_rank, is_main_process from utils.logger import log_dict_to_wandb, setup_wandb logger = logging.getLogger(__name__) def train( model, train_loaders, optimizer, tokenizer, epoch, global_step, device, scheduler, scaler, config, ): model.train() metric_logger = MetricLogger(delimiter=" ") metric_logger.add_meter("lr", SmoothedValue(window=1, fmt="{value:.6f}")) metric_logger.add_meter("temperature", SmoothedValue(window=1, fmt="{value:.4f}")) loss_names = ["loss_" + k for k, v in config.criterion.loss_weight.items() if v != 0] requires_raw_text = config.model.get("requires_raw_text", False) media_types = [loader.dataset.media_type for loader in train_loaders] for name in loss_names: for m in media_types: metric_logger.add_meter(f"{m}-{name}", SmoothedValue(window=1, fmt="{value:.4f}")) header = f"Train Epoch: [{epoch}]" log_freq = config.log_freq if config.distributed: for d in train_loaders: d.sampler.set_epoch(epoch) train_loader = MetaLoader(name2loader=dict(list(zip(media_types, train_loaders)))) model_without_ddp = model.module if config.distributed else model iterator = metric_logger.log_every(train_loader, log_freq, header) for i, (media_type, (image, text, idx)) in enumerate(iterator): image = image.to(device, non_blocking=True) idx = idx.to(device, non_blocking=True) text_input = tokenizer( text, padding="max_length", truncation=True, max_length=config.max_txt_l, return_tensors="pt", ).to(device) #with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.bfloat16): with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.float16): if requires_raw_text: loss_dict = model(image, text_input, idx=idx, raw_text=text) else: loss_dict = model(image, text_input, idx=idx) loss = sum(loss_dict.values()) #! We do not use scaler as we only involve bf16, check this optimizer.zero_grad() loss.backward() if config.optimizer.max_grad_norm > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), config.optimizer.max_grad_norm) optimizer.step() scheduler.step() # optimizer.zero_grad() # scaler.scale(loss).backward() # if config.optimizer.max_grad_norm > 0: # scaler.unscale_(optimizer) # torch.nn.utils.clip_grad_norm_(model.parameters(), config.optimizer.max_grad_norm) # scaler.step(optimizer) # scaler.update() # scheduler.step() # logging for name in loss_names: value = loss_dict[name] value = value if isinstance(value, float) else value.item() metric_logger.update(**{f"{media_type}-{name}": value}) metric_logger.update(lr=optimizer.param_groups[0]["lr"]) metric_logger.update(temperature=model_without_ddp.temp.item()) if is_main_process() and config.wandb.enable and global_step % log_freq == 0: logs = metric_logger.get_global_avg_dict() log_dict_to_wandb(logs, step=global_step, prefix="train/") global_step += 1 if config.debug and (i + 1) % 5 == 0: break # gather the stats from all processes metric_logger.synchronize_between_processes() logger.info(f"Averaged train stats: {metric_logger.global_avg()}") return global_step def main(config): if is_main_process() and config.wandb.enable: run = setup_wandb(config) logger.info(f"config: \n{config}") logger.info(f"train_file: {config.train_file}") setup_seed(config.seed + get_rank()) device = torch.device(config.device) cudnn.benchmark = True train_loaders, test_name2loaders, train_media_types = setup_dataloaders(config, mode="ret") num_steps_per_epoch = sum(len(d) for d in train_loaders) config.scheduler.num_training_steps = num_steps_per_epoch * config.scheduler.epochs config.scheduler.num_warmup_steps = num_steps_per_epoch * config.scheduler.warmup_epochs model_cls = eval(config.model.get('model_cls', 'ViCLIP')) find_unused_parameters = True if any([x in config.model.get('model_cls', 'VindLU') for x in ['VindLU_BLIP', 'VindLU_VideoCLIP']]): find_unused_parameters = False ( model, model_without_ddp, optimizer, scheduler, scaler, tokenizer, start_epoch, global_step, ) = setup_model( config, model_cls=model_cls, has_decoder=False, pretrain=False, # find_unused_parameters=True find_unused_parameters=find_unused_parameters, ) if is_main_process() and config.wandb.enable: wandb.watch(model) best = 0 best_epoch = 0 logger.info("Start " + "evaluation" if config.evaluate else "training") start_time = time.time() for epoch in range(start_epoch, config.scheduler.epochs): if not config.evaluate: global_step = train( model, train_loaders, optimizer, tokenizer, epoch, global_step, device, scheduler, scaler, config, ) if config.get('no_test', False) and not config.evaluate: save_obj = { "model": model_without_ddp.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "scaler": scaler.state_dict(), "config": config, "epoch": epoch, "global_step": global_step, } torch.save(save_obj, join(config.output_dir, "ckpt_best.pth")) best_epoch = epoch else: #with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.bfloat16): with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.float16): eval_res = {} for test_name, test_loader in test_name2loaders.items(): if test_name not in config.test_types: logger.info( f"Skip eval {test_name} split. All test_types {config.test_types}" ) continue res = evaluation_wrapper( model_without_ddp, test_loader, tokenizer, device, config, prefix=test_name ) eval_res.update(res) if is_main_process(): if config.wandb.enable: for p, v in eval_res.items(): log_dict_to_wandb(v, step=global_step, prefix=p) if config.stop_key is not None and config.stop_key in eval_res: if config.model.multimodal.enable: cur_r_mean = eval_res[config.stop_key]["r_mean"] else: cur_r_mean = eval_res[config.stop_key.replace("/", "_emb/")]["r_mean"] else: # None cur_r_mean = best + 1 # save the last as the best eval_res = pd.DataFrame(eval_res) logger.info(f"Epoch {epoch}") logger.info(f"\n{eval_res.transpose().to_string(max_cols=30)}") eval_res.to_json(join(config.output_dir, "eval_res_latest.json")) if not config.evaluate and cur_r_mean > best: save_obj = { "model": model_without_ddp.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "scaler": scaler.state_dict(), "config": config, "epoch": epoch, "global_step": global_step, } eval_file = "eval_res_best.json" eval_res.to_json(join(config.output_dir, eval_file)) torch.save(save_obj, join(config.output_dir, "ckpt_best.pth")) best = cur_r_mean best_epoch = epoch if config.evaluate: eval_file = "eval_res.json" eval_res.to_json(join(config.output_dir, eval_file)) if config.evaluate or config.debug: break dist.barrier() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) logger.info(f"Training time {total_time_str}") logger.info(f"best epoch {best_epoch} [config.stop_key {config.stop_key}]") logger.info(f"Checkpoints and Logs saved at {config.output_dir}") if is_main_process() and config.wandb.enable: run.finish() def eval_after_training(train_config): # general config for all train_config.wandb.enable = False train_config.evaluate = True train_config.pretrained_path = join(train_config.output_dir, "ckpt_best.pth") train_config.num_frames_test = train_config.num_frames train_config.inputs.video_input.num_frames_test = train_config.num_frames if train_config.get('num_frames_test_final', False): train_config.num_frames_test = train_config.num_frames_test_final train_config.batch_size = train_config.batch_size_final train_config.inputs.video_input.num_frames_test = train_config.num_frames_test_final train_config.model.vision_encoder.num_frames = train_config.num_frames_test_final eval_config = copy.deepcopy(train_config) eval_config.test_types = list(eval_config.test_file.keys()) eval_config.output_dir = join(eval_config.output_dir, f"eval_after_training") eval_config.result_dir = eval_config.output_dir if is_main_process(): os.makedirs(eval_config.output_dir, exist_ok=True) Config.dump(eval_config, os.path.join(eval_config.output_dir, "config.json")) logger.info(f"===========> START eval_after_training [{eval_config.test_types}]") main(eval_config) if __name__ == "__main__": cfg = setup_main() main(cfg) if not cfg.evaluate: eval_after_training(cfg)