|
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) |
|
|
