Spaces:
Running
Running
| import torch | |
| import yaml | |
| import time | |
| from collections import OrderedDict,namedtuple | |
| import os | |
| import sys | |
| ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| sys.path.insert(0, ROOT_DIR) | |
| from sgmnet import matcher as SGM_Model | |
| from superglue import matcher as SG_Model | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--matcher_name', type=str, default='SGM', | |
| help='number of processes.') | |
| parser.add_argument('--config_path', type=str, default='configs/cost/sgm_cost.yaml', | |
| help='number of processes.') | |
| parser.add_argument('--num_kpt', type=int, default=4000, | |
| help='keypoint number, default:100') | |
| parser.add_argument('--iter_num', type=int, default=100, | |
| help='keypoint number, default:100') | |
| def test_cost(test_data,model): | |
| with torch.no_grad(): | |
| #warm up call | |
| _=model(test_data) | |
| torch.cuda.synchronize() | |
| a=time.time() | |
| for _ in range(int(args.iter_num)): | |
| _=model(test_data) | |
| torch.cuda.synchronize() | |
| b=time.time() | |
| print('Average time per run(ms): ',(b-a)/args.iter_num*1e3) | |
| print('Peak memory(MB): ',torch.cuda.max_memory_allocated()/1e6) | |
| if __name__=='__main__': | |
| torch.backends.cudnn.benchmark=False | |
| args = parser.parse_args() | |
| with open(args.config_path, 'r') as f: | |
| model_config = yaml.load(f) | |
| model_config=namedtuple('model_config',model_config.keys())(*model_config.values()) | |
| if args.matcher_name=='SGM': | |
| model = SGM_Model(model_config) | |
| elif args.matcher_name=='SG': | |
| model = SG_Model(model_config) | |
| model.cuda(),model.eval() | |
| test_data = { | |
| 'x1':torch.rand(1,args.num_kpt,2).cuda()-0.5, | |
| 'x2':torch.rand(1,args.num_kpt,2).cuda()-0.5, | |
| 'desc1': torch.rand(1,args.num_kpt,128).cuda(), | |
| 'desc2': torch.rand(1,args.num_kpt,128).cuda() | |
| } | |
| test_cost(test_data,model) | |