| 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(): |
| |
| _ = 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) |
|
|
