import os import sys import argparse import random import datetime import matplotlib as mpl mpl.use('Agg') import numpy as np import torch import torch.nn.parallel import torch.backends.cudnn as cudnn import torch.distributed as dist import torch.optim import torch.utils.data.distributed from torch.utils.data import DataLoader from torchvision.transforms import Compose, ToTensor, Normalize import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP import torch.utils.data.distributed #import apex.amp as amp from torch.cuda.amp import autocast as autocast from model.model_sbert_gref import Model_CL from model.model_sbert_gref import * from engine.engine_gref_sbert_oiou import * from dataset.data_loader_test import * from utils.losses import * from utils.parsing_metrics import * from utils.utils import * from utils.checkpoint import load_pretrain, load_resume from utils.logger import setup_logger def get_args(): parser = argparse.ArgumentParser(description='Dataloader test') parser.add_argument('--gpu', default='2', help='gpu id') parser.add_argument('--ngpu', default=2, type=int, help='gpu num') parser.add_argument('--workers', default=4, type=int, help='num workers for data loading') parser.add_argument('--seed', default=0, type=int, help='random seed') parser.add_argument('--clip_model', default='ViT-B/16', type=str, help='clip model RN50 RN101 ViT-B/32') parser.add_argument('--nb_epoch', default=32, type=int, help='training epoch') parser.add_argument('--lr', default=0.000025, type=float, help='batch size 16 learning rate') parser.add_argument('--power', default=0.1, type=float, help='lr poly power') parser.add_argument('--steps', default=[15, 28], type=list, help='in which step lr decay by power') parser.add_argument('--batch_size', default=16, type=int, help='batch size') parser.add_argument('--size', default=416, type=int, help='image size') parser.add_argument('--dataset', default='refcoco', type=str, help='refcoco/refcoco+/refcocog/grefcoco') parser.add_argument('--num_query', default=16, type=int, help='the number of query') parser.add_argument('--w_seg', default=0.1, type=float, help='weight of the seg loss') parser.add_argument('--w_coord', default=5, type=float, help='weight of the reg loss') parser.add_argument('--tunelang', dest='tunelang', default=True, action='store_true', help='if finetune language model') parser.add_argument('--anchor_imsize', default=416, type=int, help='scale used to calculate anchors defined in model cfg file') parser.add_argument('--data_root', type=str, default='./ln_data', help='path to ReferIt splits data folder') parser.add_argument('--split_root', type=str, default='./data', help='location of pre-parsed dataset info') parser.add_argument('--time', default=15, type=int, help='maximum time steps (lang length) per batch') parser.add_argument('--log_dir', type=str, default='./logs', help='path to ReferIt splits data folder') parser.add_argument('--fusion_dim', default=768, type=int, help='fusion module embedding dimensions') parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--pretrain', default='', type=str, metavar='PATH', help='pretrain support load state_dict that are not identical, while have no loss saved as resume') parser.add_argument('--print_freq', '-p', default=100, type=int, metavar='N', help='print frequency (default: 1e3)') parser.add_argument('--savename', default='default', type=str, help='Name head for saved model') parser.add_argument('--seg_thresh', default=0.35, type=float, help='seg score above this value means foreground') parser.add_argument('--seg_out_stride', default=2, type=int, help='the seg out stride') parser.add_argument('--best_iou', default=-float('Inf'), type=int, help='the best accu') parser.add_argument('--fuse_mode', default='coarse') global args, anchors_full, writer, logger args = parser.parse_args() args.gsize = 32 args.date = datetime.datetime.now().strftime('%Y%m%d') if args.savename=='default': args.savename = 'model_v1_%s_batch%d_%s'%(args.dataset, args.batch_size, args.date) os.makedirs(args.log_dir, exist_ok=True) args.lr = args.lr * (args.batch_size * args.ngpu // 16) print('----------------------------------------------------------------------') print(sys.argv[0]) print(args) print('----------------------------------------------------------------------') return args def main(args): os.environ['MASTER_ADDR'] = 'localhost' # os.environ['MASTER_PORT'] = '12367' if(torch.cuda.is_available()): n_gpus = torch.cuda.device_count() print("Running DDP with {} GPUs".format(n_gpus)) mp.spawn(run, nprocs=n_gpus, args=(n_gpus, args,)) else: print("Please use GPU for training") def run(rank, n_gpus, args): dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank) torch.cuda.set_device(rank) ## fix seed cudnn.benchmark = False cudnn.deterministic = True random.seed(args.seed) np.random.seed(args.seed+1) torch.manual_seed(args.seed+2) torch.cuda.manual_seed_all(args.seed+3) ## save logs logger = setup_logger(output=os.path.join(args.log_dir, args.savename), distributed_rank=rank, color=False, name="model-v1") logger.info(str(sys.argv)) logger.info(str(args)) input_transform = Compose([ ToTensor(), Normalize( mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711] ) ]) val_dataset = ReferDataset(data_root=args.data_root, dataset=args.dataset, split_root=args.split_root, split='val', imsize = args.size, transform=input_transform, max_query_len=args.time) val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, pin_memory=True, drop_last=True, num_workers=args.workers) if args.dataset == 'refcocog_u': test_dataset = ReferDataset(data_root=args.data_root, dataset=args.dataset, split_root=args.split_root, split='test', imsize = args.size, transform=input_transform, max_query_len=args.time) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, pin_memory=True, drop_last=True, num_workers=args.workers) elif args.dataset == 'refcocog_g': pass else: testA_dataset = ReferDataset(data_root=args.data_root, dataset=args.dataset, split_root=args.split_root, split='testA', imsize = args.size, transform=input_transform, max_query_len=args.time) testB_dataset = ReferDataset(data_root=args.data_root, dataset=args.dataset, split_root=args.split_root, split='testB', imsize = args.size, transform=input_transform, max_query_len=args.time) testA_loader = DataLoader(testA_dataset, batch_size=1, shuffle=False, pin_memory=True, drop_last=True, num_workers=args.workers) testB_loader = DataLoader(testB_dataset, batch_size=1, shuffle=False, pin_memory=True, drop_last=True, num_workers=args.workers) ## Model model = Model_CL(clip_model=args.clip_model, tunelang=args.tunelang, num_query=args.num_query, fusion_dim=args.fusion_dim, fuse_mode=args.fuse_mode, use_projections=True).cuda(rank) model = DDP(model, device_ids=[rank], find_unused_parameters=True) model_without_ddp = model.module args.start_epoch = 0 if args.pretrain and os.path.isfile(args.pretrain): model=load_pretrain(model, args, logger, rank) model.to(rank) visu_param = [param for name, param in model_without_ddp.named_parameters() if 'visumodel' in name] text_param = [param for name, param in model_without_ddp.named_parameters() if 'textmodel' in name] rest_param = [param for name, param in model_without_ddp.named_parameters() if 'textmodel' not in name and 'visumodel' not in name] ## optimizer; adam default if args.tunelang: optimizer = torch.optim.Adam([{'params': rest_param, 'lr': args.lr}, {'params': visu_param, 'lr': args.lr / 10.}, {'params': text_param, 'lr': args.lr / 10.}]) else: optimizer = torch.optim.Adam([{'params': rest_param}, {'params': visu_param, 'lr': args.lr / 10.}], lr=args.lr) best_miou_seg = -float('Inf') if args.resume: model = load_resume(model, optimizer, args, logger, rank) model.to(rank) best_miou_seg = args.best_iou print(best_miou_seg) if args.dataset == 'refcocog_u': print('\nTest testing:') miou_seg, oiou_seg, prec = validate_epoch(args, test_loader, model, logger, 'test') elif args.dataset == 'refcocog_g': pass else: print('\nTestA testing:') miou_seg, oiou_seg, prec = validate_epoch(args, testA_loader, model, logger, 'testA') print('\nTestB testing:') miou_seg,oiou_seg, prec = validate_epoch(args, testB_loader, model, logger, 'testB') miou_seg, oiou_seg, prec = validate_epoch(args, val_loader, model, logger, 'val') if __name__ == "__main__": args = get_args() main(args)