-
Notifications
You must be signed in to change notification settings - Fork 11
/
test_net.py
60 lines (48 loc) · 1.67 KB
/
test_net.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import os
import random
import argparse
import torch
import numpy as np
from benchmark.config import cfg
from benchmark.models import build_model
from benchmark.evaluation import ROxford5kEvaluator, RParis6kEvaluator, InstreEvaluator, GLDRetrievalEvaluator
from benchmark.utils.logger import setup_logger
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.backends.cudnn.deterministic = True
def test():
logger = setup_logger("insclr")
device = torch.device(cfg.DEVICE)
model = build_model(cfg, logger)
logger.info(model)
model.to(device)
if len(os.environ["CUDA_VISIBLE_DEVICES"]) > 1:
model = torch.nn.DataParallel(model)
logger.info(cfg)
# evaluator
evaluators = []
# select evaluator according to your need
# evaluators.append(GLDRetrievalEvaluator("datasets/google-landmark"))
# evaluators.append(ROxford5kEvaluator("datasets/roxford5k", 1024, cfg.TEST.SCALES))
# evaluators.append(RParis6kEvaluator("datasets/rparis6k", 1024, cfg.TEST.SCALES))
# evaluators.append(InstreEvaluator("datasets/instre", 256, cfg.TEST.SCALES, pca=True))
for evaluator in evaluators:
logger.info(evaluator)
for evaluator in evaluators:
msg, _ = evaluator.evaluate(model, device)
logger.info(msg)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--cfg", type=str)
parser.add_argument(
"opts",
help="modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg.merge_from_file(args.cfg)
cfg.merge_from_list(args.opts)
test()