-
Notifications
You must be signed in to change notification settings - Fork 98
/
test.py
36 lines (28 loc) · 1.08 KB
/
test.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
import torch, os
from utils.common import merge_config, get_model
from evaluation.eval_wrapper import eval_lane
import torch
if __name__ == "__main__":
torch.backends.cudnn.benchmark = True
args, cfg = merge_config()
distributed = False
if 'WORLD_SIZE' in os.environ:
distributed = int(os.environ['WORLD_SIZE']) > 1
cfg.distributed = distributed
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
net = get_model(cfg)
state_dict = torch.load(cfg.test_model, map_location = 'cpu')['model']
compatible_state_dict = {}
for k, v in state_dict.items():
if 'module.' in k:
compatible_state_dict[k[7:]] = v
else:
compatible_state_dict[k] = v
net.load_state_dict(compatible_state_dict, strict = True)
if distributed:
net = torch.nn.parallel.DistributedDataParallel(net, device_ids = [args.local_rank])
if not os.path.exists(cfg.test_work_dir):
os.mkdir(cfg.test_work_dir)
eval_lane(net, cfg)