if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
dist_print('start testing...')
assert cfg.backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']
if cfg.dataset == 'CULane':
cls_num_per_lane = 18
elif cfg.dataset == 'Tusimple':
cls_num_per_lane = 56
else:
raise NotImplementedError
net = parsingNet(pretrained = False, backbone=cfg.backbone,cls_dim = (cfg.griding_num+1,cls_num_per_lane,4),
use_aux=False).cuda() # we dont need auxiliary segmentation in testing
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 = False)
if distributed:
net = torch.nn.parallel.DistributedDataParallel(net, device_ids = [args.local_rank])
eval_lane(net, cfg.dataset, cfg.data_root, cfg.test_work_dir, cfg.griding_num, False, distributed)
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
dist_print('start testing...')
net = E2ENet(Channels=96,
nums_lane=4,
culomn_channels=cfg.griding_num,
row_channels=cfg.row_num,
initialed=True).cuda()
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=False)
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.dataset, cfg.data_root, cfg.test_work_dir, distributed,
cfg)
cls_num_per_lane = 18
elif cfg['dataset']['name'] == 'Tusimple':
cls_num_per_lane = 56
else:
raise NotImplementedError
net = parsingNet(network=cfg['network'], datasets=cfg['dataset']).cuda()
state_dict = torch.load(cfg['test']['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=False)
# if distributed:
# net = torch.nn.parallel.DistributedDataParallel(net, device_ids = [args.local_rank])
if not os.path.exists(cfg['test']['test_work_dir']):
os.mkdir(cfg['test']['test_work_dir'])
dataset = cfg['dataset']['name']
data_root = cfg['dataset']['data_root']
test_work_dir = cfg['test']['test_work_dir']
griding_num = cfg['dataset']['griding_num']
eval_lane(net, dataset, data_root, test_work_dir, griding_num, True,
distributed)