device = torch.device("cuda") else: device = torch.device("cpu") # input size input_size = [args.input_size, args.input_size] # load net if args.version == 'centernet': from models.centernet import CenterNet net = CenterNet(device, input_size=input_size, num_classes=num_classes, backbone=args.backbone, use_nms=args.use_nms) # load net net.load_state_dict(torch.load(args.trained_model, map_location='cuda')) net.eval() print('Finished loading model!') net = net.to(device) # evaluation with torch.no_grad(): if args.dataset == 'voc': voc_test(net, device, input_size) elif args.dataset == 'coco-val': coco_test(net, device, input_size, test=False) elif args.dataset == 'coco-test': coco_test(net, device, input_size, test=True)
print('ap50_95 : ', ap50_95) if __name__ == '__main__': global cfg cfg = coco_cfg if args.cuda: print('use cuda') torch.backends.cudnn.benchmark = True device = torch.device("cuda") else: device = torch.device("cpu") if args.version == 'centernet': from models.centernet import CenterNet model = CenterNet(device, input_size=cfg['min_dim'], num_classes=args.num_classes) else: print('Unknown Version !!!') exit() # load model model.load_state_dict(torch.load(args.trained_model, map_location=device)) model.eval().to(device) print('Finished loading model!') test(model, device)