dianecy commited on
Commit
0218947
·
verified ·
1 Parent(s): 931add2

Upload ./ASDA/test_proj.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. ASDA/test_proj.py +236 -0
ASDA/test_proj.py ADDED
@@ -0,0 +1,236 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import sys
3
+ import argparse
4
+ import random
5
+ import datetime
6
+ import matplotlib as mpl
7
+ mpl.use('Agg')
8
+
9
+ import numpy as np
10
+ import torch
11
+ import torch.nn.parallel
12
+ import torch.backends.cudnn as cudnn
13
+ import torch.distributed as dist
14
+ import torch.optim
15
+ import torch.utils.data.distributed
16
+ from torch.utils.data import DataLoader
17
+ from torchvision.transforms import Compose, ToTensor, Normalize
18
+
19
+ import torch.distributed as dist
20
+ import torch.multiprocessing as mp
21
+ from torch.nn.parallel import DistributedDataParallel as DDP
22
+ import torch.utils.data.distributed
23
+
24
+ #import apex.amp as amp
25
+ from torch.cuda.amp import autocast as autocast
26
+
27
+ from model.model_sbert_gref import Model_CL
28
+ from model.model_sbert_gref import *
29
+ from engine.engine_gref_sbert_oiou import *
30
+
31
+ from dataset.data_loader_test import *
32
+ from utils.losses import *
33
+ from utils.parsing_metrics import *
34
+ from utils.utils import *
35
+ from utils.checkpoint import load_pretrain, load_resume
36
+ from utils.logger import setup_logger
37
+
38
+ def get_args():
39
+ parser = argparse.ArgumentParser(description='Dataloader test')
40
+ parser.add_argument('--gpu', default='2', help='gpu id')
41
+ parser.add_argument('--ngpu', default=2, type=int, help='gpu num')
42
+ parser.add_argument('--workers', default=4, type=int, help='num workers for data loading')
43
+ parser.add_argument('--seed', default=0, type=int, help='random seed')
44
+
45
+ parser.add_argument('--clip_model', default='ViT-B/16', type=str, help='clip model RN50 RN101 ViT-B/32')
46
+ parser.add_argument('--nb_epoch', default=32, type=int, help='training epoch')
47
+ parser.add_argument('--lr', default=0.000025, type=float, help='batch size 16 learning rate')
48
+ parser.add_argument('--power', default=0.1, type=float, help='lr poly power')
49
+ parser.add_argument('--steps', default=[15, 28], type=list, help='in which step lr decay by power')
50
+ parser.add_argument('--batch_size', default=16, type=int, help='batch size')
51
+ parser.add_argument('--size', default=416, type=int, help='image size')
52
+ parser.add_argument('--dataset', default='refcoco', type=str,
53
+ help='refcoco/refcoco+/refcocog/grefcoco')
54
+
55
+ parser.add_argument('--num_query', default=16, type=int, help='the number of query')
56
+ parser.add_argument('--w_seg', default=0.1, type=float, help='weight of the seg loss')
57
+ parser.add_argument('--w_coord', default=5, type=float, help='weight of the reg loss')
58
+ parser.add_argument('--tunelang', dest='tunelang', default=True, action='store_true', help='if finetune language model')
59
+ parser.add_argument('--anchor_imsize', default=416, type=int,
60
+ help='scale used to calculate anchors defined in model cfg file')
61
+ parser.add_argument('--data_root', type=str, default='./ln_data',
62
+ help='path to ReferIt splits data folder')
63
+ parser.add_argument('--split_root', type=str, default='./data',
64
+ help='location of pre-parsed dataset info')
65
+ parser.add_argument('--time', default=15, type=int,
66
+ help='maximum time steps (lang length) per batch')
67
+ parser.add_argument('--log_dir', type=str, default='./logs',
68
+ help='path to ReferIt splits data folder')
69
+
70
+ parser.add_argument('--fusion_dim', default=768, type=int,
71
+ help='fusion module embedding dimensions')
72
+ parser.add_argument('--resume', default='', type=str, metavar='PATH',
73
+ help='path to latest checkpoint (default: none)')
74
+ parser.add_argument('--pretrain', default='', type=str, metavar='PATH',
75
+ help='pretrain support load state_dict that are not identical, while have no loss saved as resume')
76
+ parser.add_argument('--print_freq', '-p', default=100, type=int,
77
+ metavar='N', help='print frequency (default: 1e3)')
78
+ parser.add_argument('--savename', default='default', type=str, help='Name head for saved model')
79
+
80
+ parser.add_argument('--seg_thresh', default=0.35, type=float, help='seg score above this value means foreground')
81
+ parser.add_argument('--seg_out_stride', default=2, type=int, help='the seg out stride')
82
+ parser.add_argument('--best_iou', default=-float('Inf'), type=int, help='the best accu')
83
+ parser.add_argument('--fuse_mode', default='coarse')
84
+
85
+
86
+
87
+ global args, anchors_full, writer, logger
88
+ args = parser.parse_args()
89
+ args.gsize = 32
90
+ args.date = datetime.datetime.now().strftime('%Y%m%d')
91
+ if args.savename=='default':
92
+ args.savename = 'model_v1_%s_batch%d_%s'%(args.dataset, args.batch_size, args.date)
93
+ os.makedirs(args.log_dir, exist_ok=True)
94
+ args.lr = args.lr * (args.batch_size * args.ngpu // 16)
95
+
96
+ print('----------------------------------------------------------------------')
97
+ print(sys.argv[0])
98
+ print(args)
99
+ print('----------------------------------------------------------------------')
100
+
101
+ return args
102
+
103
+ def main(args):
104
+ os.environ['MASTER_ADDR'] = 'localhost'
105
+ # os.environ['MASTER_PORT'] = '12367'
106
+
107
+ if(torch.cuda.is_available()):
108
+ n_gpus = torch.cuda.device_count()
109
+ print("Running DDP with {} GPUs".format(n_gpus))
110
+ mp.spawn(run, nprocs=n_gpus, args=(n_gpus, args,))
111
+ else:
112
+ print("Please use GPU for training")
113
+
114
+ def run(rank, n_gpus, args):
115
+ dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
116
+ torch.cuda.set_device(rank)
117
+
118
+ ## fix seed
119
+ cudnn.benchmark = False
120
+ cudnn.deterministic = True
121
+ random.seed(args.seed)
122
+ np.random.seed(args.seed+1)
123
+ torch.manual_seed(args.seed+2)
124
+ torch.cuda.manual_seed_all(args.seed+3)
125
+
126
+ ## save logs
127
+ logger = setup_logger(output=os.path.join(args.log_dir, args.savename), distributed_rank=rank, color=False, name="model-v1")
128
+ logger.info(str(sys.argv))
129
+ logger.info(str(args))
130
+
131
+ input_transform = Compose([
132
+ ToTensor(),
133
+ Normalize(
134
+ mean=[0.48145466, 0.4578275, 0.40821073],
135
+ std=[0.26862954, 0.26130258, 0.27577711]
136
+ )
137
+ ])
138
+
139
+
140
+ val_dataset = ReferDataset(data_root=args.data_root,
141
+ dataset=args.dataset,
142
+ split_root=args.split_root,
143
+ split='val',
144
+ imsize = args.size,
145
+ transform=input_transform,
146
+ max_query_len=args.time)
147
+
148
+
149
+ val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False,
150
+ pin_memory=True, drop_last=True, num_workers=args.workers)
151
+
152
+ if args.dataset == 'refcocog_u':
153
+ test_dataset = ReferDataset(data_root=args.data_root,
154
+ dataset=args.dataset,
155
+ split_root=args.split_root,
156
+ split='test',
157
+ imsize = args.size,
158
+ transform=input_transform,
159
+ max_query_len=args.time)
160
+
161
+ test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False,
162
+ pin_memory=True, drop_last=True, num_workers=args.workers)
163
+ elif args.dataset == 'refcocog_g':
164
+ pass
165
+ else:
166
+ testA_dataset = ReferDataset(data_root=args.data_root,
167
+ dataset=args.dataset,
168
+ split_root=args.split_root,
169
+ split='testA',
170
+ imsize = args.size,
171
+ transform=input_transform,
172
+ max_query_len=args.time)
173
+ testB_dataset = ReferDataset(data_root=args.data_root,
174
+ dataset=args.dataset,
175
+ split_root=args.split_root,
176
+ split='testB',
177
+ imsize = args.size,
178
+ transform=input_transform,
179
+ max_query_len=args.time)
180
+
181
+
182
+ testA_loader = DataLoader(testA_dataset, batch_size=1, shuffle=False,
183
+ pin_memory=True, drop_last=True, num_workers=args.workers)
184
+ testB_loader = DataLoader(testB_dataset, batch_size=1, shuffle=False,
185
+ pin_memory=True, drop_last=True, num_workers=args.workers)
186
+
187
+
188
+ ## Model
189
+ 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)
190
+ model = DDP(model, device_ids=[rank], find_unused_parameters=True)
191
+ model_without_ddp = model.module
192
+
193
+ args.start_epoch = 0
194
+ if args.pretrain and os.path.isfile(args.pretrain):
195
+ model=load_pretrain(model, args, logger, rank)
196
+ model.to(rank)
197
+
198
+ visu_param = [param for name, param in model_without_ddp.named_parameters() if 'visumodel' in name]
199
+ text_param = [param for name, param in model_without_ddp.named_parameters() if 'textmodel' in name]
200
+ rest_param = [param for name, param in model_without_ddp.named_parameters() if 'textmodel' not in name and 'visumodel' not in name]
201
+
202
+
203
+ ## optimizer; adam default
204
+ if args.tunelang:
205
+ optimizer = torch.optim.Adam([{'params': rest_param, 'lr': args.lr},
206
+ {'params': visu_param, 'lr': args.lr / 10.},
207
+ {'params': text_param, 'lr': args.lr / 10.}])
208
+ else:
209
+ optimizer = torch.optim.Adam([{'params': rest_param},
210
+ {'params': visu_param, 'lr': args.lr / 10.}], lr=args.lr)
211
+
212
+
213
+
214
+ best_miou_seg = -float('Inf')
215
+ if args.resume:
216
+ model = load_resume(model, optimizer, args, logger, rank)
217
+ model.to(rank)
218
+ best_miou_seg = args.best_iou
219
+ print(best_miou_seg)
220
+
221
+ if args.dataset == 'refcocog_u':
222
+ print('\nTest testing:')
223
+ miou_seg, oiou_seg, prec = validate_epoch(args, test_loader, model, logger, 'test')
224
+
225
+ elif args.dataset == 'refcocog_g':
226
+ pass
227
+ else:
228
+ print('\nTestA testing:')
229
+ miou_seg, oiou_seg, prec = validate_epoch(args, testA_loader, model, logger, 'testA')
230
+ print('\nTestB testing:')
231
+ miou_seg,oiou_seg, prec = validate_epoch(args, testB_loader, model, logger, 'testB')
232
+ miou_seg, oiou_seg, prec = validate_epoch(args, val_loader, model, logger, 'val')
233
+
234
+ if __name__ == "__main__":
235
+ args = get_args()
236
+ main(args)