import torch #from utils import concat_all_gather, is_dist_avail_and_initialized, accuracy #the original concat_all_gather is abandoned because of no gradient backward from utils import is_dist_avail_and_initialized, accuracy import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist from tqdm import tqdm import sys sys.path.append("..") from sharegpt4v import share4v_val_dataset, share4v_train_dataset from model import longclip from torch.utils.data.distributed import DistributedSampler from scheduler import cosine_lr import argparse import os import subprocess import collections import torch.optim as optim from torch.utils.tensorboard import SummaryWriter import numpy as np from datetime import datetime from torch.cuda.amp import GradScaler # import warnings # warnings.filterwarnings("ignore") class CLIP_Clean_Train(): def __init__(self, rank,local_rank,args): self.rank=rank self.local_rank = local_rank self.base_model = args.base_model self.model, _ = longclip.load_from_clip(self.base_model, device='cpu',download_root=args.download_root) self.model.train() self.model.logit_scale = torch.nn.Parameter(torch.ones([]) * args.log_scale) self.model = self.model.cuda() self.batch_size = args.batch_size self.num_epoch = args.epochs self.lr = args.lr self.weight_decay = args.weight_decay self.warmup_length = args.warmup_length if args.exp_name == "auto": self.logdir = f"longclip/lr={args.lr}_wd={args.weight_decay}_wl={args.warmup_length}_logs={args.log_scale}_64xb" else: self.logdir = args.exp_name self.ckptdir = self.logdir + "/ckpt/" os.makedirs(self.ckptdir, exist_ok=True) self.writer = SummaryWriter(self.logdir) self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[local_rank]) self.optimizer = optim.AdamW(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.scaler =GradScaler() # move PCA to CLIP class # def PCA(self, input_tensor, PCA_dim): # mean = torch.mean(input_tensor, dim=0) # X_centered = input_tensor - mean.unsqueeze(0) # X_centered = X_centered.float() # cov_matrix = torch.mm(X_centered.T, X_centered) # eigenvalues, eigenvectors = torch.linalg.eig(cov_matrix) # eigenvalues = eigenvalues.float() # eigenvectors = eigenvectors.float() # sorted_indices = torch.argsort(eigenvalues, descending=True) # eigenvectors = eigenvectors[:, sorted_indices] # principal_components = eigenvectors[:, :PCA_dim] # X_transformed = torch.mm(X_centered, principal_components) # X_reversed = torch.mm(X_transformed, principal_components.T) # X_reversed += mean # return X_reversed #rewrite forward in CLIP class to take place inference function, such that DDP will be effective # def inference(self, images, texts): # image_features = self.model.module.encode_image(images) # image_features = image_features / image_features.norm(dim=-1, keepdim=True) # text_features = self.model.module.encode_text(texts) # text_features = text_features / text_features.norm(dim=-1, keepdim=True) # image_feat_all = concat_all_gather(image_features) # text_feat_all = concat_all_gather(text_features) # sim_i2t = torch.matmul(image_features, text_feat_all.T) # sim_t2i = torch.matmul(image_feat_all, text_features.T) # sim_t2i = sim_t2i.T # sim_i2t = self.model.logit_scale.exp() * sim_i2t # sim_t2i = self.model.logit_scale.exp() * sim_t2i # if is_dist_avail_and_initialized(): # rank = dist.get_rank() # else: # rank = 0 # bs = images.size(0) # targets = torch.linspace(rank * bs, rank * bs + bs - 1, bs, dtype=int).to( # images.device # ) # loss_itc = ( # F.cross_entropy(sim_i2t, targets, label_smoothing=0.1) # + F.cross_entropy(sim_t2i, targets, label_smoothing=0.1) # ) / 2 # return loss_itc # def inference_short(self, images, texts): # image_features = self.model.module.encode_image(images) # image_features = image_features / image_features.norm(dim=-1, keepdim=True) # image_features = self.PCA(image_features, 32) # text_features = self.model.module.encode_text(texts) # text_features = text_features / text_features.norm(dim=-1, keepdim=True) # image_feat_all = concat_all_gather(image_features) # text_feat_all = concat_all_gather(text_features) # sim_i2t = torch.matmul(image_features, text_feat_all.T) # sim_t2i = torch.matmul(image_feat_all, text_features.T) # sim_t2i = sim_t2i.T # sim_i2t = self.model.logit_scale.exp() * sim_i2t # sim_t2i = self.model.logit_scale.exp() * sim_t2i # if is_dist_avail_and_initialized(): # rank = dist.get_rank() # else: # rank = 0 # bs = images.size(0) # targets = torch.linspace(rank * bs, rank * bs + bs - 1, bs, dtype=int).to( # images.device # ) # loss_itc = ( # F.cross_entropy(sim_i2t, targets, label_smoothing=0.1) # + F.cross_entropy(sim_t2i, targets, label_smoothing=0.1) # ) / 2 # return loss_itc def train_epoch(self, dataloader, epoch, start_iter=0): running_loss = 0.0 running_loss_short = 0.0 #rank = torch.distributed.get_rank() num_batches_per_epoch = len(dataloader) for i, (images, texts, short_text) in enumerate(tqdm(dataloader, disable=(self.rank != 0))): step = num_batches_per_epoch * epoch + i if step < start_iter: continue #images = images.cuda() #images_short = images.clone() texts = longclip.tokenize(texts, truncate=True).cuda() short_text = longclip.tokenize(short_text, truncate=True).cuda() self.scheduler(step) self.optimizer.zero_grad() with torch.cuda.amp.autocast(): loss_long,loss_short = self.model(images, texts,short_text,self.rank) # try: # loss_short = 0.1 * self.inference_short(images_short, short_text) # loss.backward() # loss_short.backward() # except: # # SVD may encounter infs, very rare occasion. # loss.backward() loss=loss_long+loss_short self.scaler.scale(loss).backward() self.scaler.step(self.optimizer) self.scaler.update() # ToDo: revise the report part # running_loss += loss.item() # running_loss_short += loss_short.item() # batch_num = i # loss = running_loss # running_loss = 0.0 # loss_short = running_loss_short # running_loss_short = 0.0 # loss = torch.tensor(loss).cuda() # dist.all_reduce(loss) # loss = loss.item() / torch.distributed.get_world_size() # loss_short = torch.tensor(loss_short).cuda() # dist.all_reduce(loss_short) # loss_short = loss_short.item() / torch.distributed.get_world_size() # rank = torch.distributed.get_rank() # if step % 100 == 0: # if rank == 0: # self.writer.add_scalar("hyper/lr", self.optimizer.param_groups[0]['lr'], step) # self.writer.add_scalar("logit_scale/train", self.model.logit_scale.item(), step) # print("=====================================") # print(f"train lr step {step}: {self.optimizer.param_groups[0]['lr']}") # print(f"train logit_scale step {step}: {self.model.logit_scale.item()}") # print(f"train loss step {step}: {loss}") # print(f"train loss short step {step}: {loss_short}") # print("=====================================") # self.writer.add_scalar("Loss/train", loss + loss_short, step) # with torch.no_grad(): # self.model.eval() # self.test(epoch = epoch) # self.model.train() # return running_loss / batch_num @torch.no_grad() def test_epoch(self, dataloader): temp_corr_dict = dict() rank = torch.distributed.get_rank() for id, (images, text) in enumerate(tqdm(dataloader, disable=(rank != 0))): images = images.cuda() image_features = self.model.module.encode_image(images) image_features = image_features / image_features.norm(dim=-1, keepdim=True) text = longclip.tokenize(text, truncate=True).cuda() text_feature = self.model.module.encode_text(text) text_feature /= text_feature.norm(dim=-1, keepdim=True) i = 0 correct = 0 total = 0 for i in range(text_feature.shape[0]): text = text_feature[i] sim = text @ image_features.T sim = sim.squeeze() correct_i = torch.argmax(sim) if i==correct_i: correct = correct + 1 total = total + 1 return correct/total def test(self, epoch=0): rank = torch.distributed.get_rank() if rank == 0: self.model.eval() testset = share4v_val_dataset() testloader = torch.utils.data.DataLoader(testset, batch_size=1000, num_workers=32, pin_memory=True) with torch.no_grad(): acc = self.test_epoch(testloader) print("=====================================") print(f"test mean of share4v retrieval: {acc}") print("=====================================") return def train(self, resume=False, warmup_length=200): trainset = share4v_train_dataset() train_sampler = DistributedSampler(dataset=trainset, shuffle=True) train_loader = torch.utils.data.DataLoader(trainset, batch_size=self.batch_size, sampler=train_sampler, num_workers=32, pin_memory=True) self.scheduler = cosine_lr(self.optimizer, base_lr=self.lr, warmup_length=warmup_length, steps=self.num_epoch * len(train_loader)) start_epoch = 0 resume_iter = 0 for epoch in range(start_epoch, self.num_epoch): self.train_epoch(train_loader, epoch, start_iter=resume_iter) if self.rank == 0: name = "longclip.pt" now = datetime.now() formatted_date = now.strftime("%m-%d--%H_%M_%S_") #torch.distributed.barrier() torch.save(self.model.module.state_dict(), './checkpoints/'+str(self.rank)+formatted_date+name) # print("=====================================") # print(f"loss after training epoch: {epoch}") # print("=====================================") # if epoch == self.num_epoch - 1: # if self.base_model == "ViT-B/16": # name = 'longclip-B.pt' # elif self.base_model == "ViT-L/14": # name = 'longclip-L.pt' # else: # name = "longclip-others.pt" # torch.save(self.model.module.state_dict(), name) def setup_distributed(backend="nccl", port=None): """Initialize distributed training environment. support both slurm and torch.distributed.launch see torch.distributed.init_process_group() for more details """ num_gpus = torch.cuda.device_count() if "SLURM_JOB_ID" in os.environ: rank = int(os.environ["SLURM_PROCID"]) world_size = int(os.environ["SLURM_NTASKS"]) node_list = os.environ["SLURM_NODELIST"] addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1") # specify master port if port is not None: os.environ["MASTER_PORT"] = str(port) elif "MASTER_PORT" not in os.environ: os.environ["MASTER_PORT"] = "29522" if "MASTER_ADDR" not in os.environ: os.environ["MASTER_ADDR"] = addr os.environ["WORLD_SIZE"] = str(world_size) os.environ["LOCAL_RANK"] = str(rank % num_gpus) os.environ["RANK"] = str(rank) else: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) torch.cuda.set_device(rank % num_gpus) dist.init_process_group( backend=backend, world_size=world_size, rank=rank, ) torch.cuda.set_device(device=f'cuda:{rank % num_gpus}') return rank, rank % num_gpus if __name__ == "__main__": parser = argparse.ArgumentParser(description='params') parser.add_argument('--lr', default=1e-6, type=float, help='lr.') parser.add_argument('--weight_decay', default=1e-2, type=float, help='wd.') parser.add_argument('--log_scale', default=4.6052, type=float, help='clip temperature log scale.') parser.add_argument("--exp_name", default="auto", type=str, help="specify experiment name.") parser.add_argument("--warmup_length", default=200, type=int, help="warmup_length.") parser.add_argument("--base_model", default="ViT-L/14", help="CLIP Base Model") parser.add_argument( "--batch-size", type=int, default=128, help="Batch size per gpu."#112 ) parser.add_argument( "--epochs", type=int, default=2, help="Number of epochs to train for." ) parser.add_argument( "--resume", default=False, action='store_true', help="resume training from checkpoint." ) parser.add_argument("--download-root", default=None, help="CLIP Base Model download root") args = parser.parse_args() rank,local_rank = setup_distributed() print("DDP Done") trainer = CLIP_Clean_Train( rank=rank, local_rank=local_rank, args=args ) trainer.train(resume=args.resume, warmup_length=args.warmup_length)