def __init__(self, model_path, long_side, network): torch.set_grad_enabled(False) if network == 'mobilenet': self.cfg = cfg_mnet net = RetinaFace(cfg=self.cfg, phase='test') elif network == 'slim': self.cfg = cfg_slim net = Slim(cfg=self.cfg, phase='test') elif network == 'RFB': self.cfg = cfg_rfb net = RFB(cfg=self.cfg, phase='test') else: print("not supported network!!") exit(0) self.net = load_model(net, model_path, True) self.net.eval() print("Finished loading model!") cudnn.benchmark = True self.device = torch.device("cpu") self.net = self.net.to(self.device) self.long_side = long_side
if __name__ == '__main__': torch.set_grad_enabled(False) cfg = None net = None if args.network == "mobile0.25": cfg = cfg_mnet net = RetinaFace(cfg = cfg, phase = 'test') elif args.network == "slim": cfg = cfg_slim net = Slim(cfg = cfg, phase = 'test') elif args.network == "RFB": cfg = cfg_rfb net = RFB(cfg = cfg, phase = 'test') else: print("Don't support network!") exit(0) net = load_model(net, args.trained_model, args.cpu) net.eval() print('Finished loading model!') print(net) cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) cap = cv2.VideoCapture(args.video_path) # testing begin
args = parser.parse_args() if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) cfg = None net = None if args.network == "mobile0.25": cfg = cfg_mnet net = RetinaFace(cfg=cfg) elif args.network == "slim": cfg = cfg_slim net = Slim(cfg=cfg) elif args.network == "RFB": cfg = cfg_rfb net = RFB(cfg=cfg) elif args.network == "efficientdet": cfg = cfg_efficientdet net = EfficientDet(cfg=cfg) else: print("Don't support network!") exit(0) print("Printing net...") #print(net) rgb_mean = (104, 117, 123) # bgr order num_classes = 2 img_dim = cfg['image_size'] num_gpu = cfg['ngpu'] batch_size = cfg['batch_size']
# # torch.set_grad_enabled(False) # x_onnx=np.float32(x) # for i in range(1000): # t1_onnx = time.time() # y_onnx = session.run([], {input_name: x_onnx}) # t1_onnxend = time.time()- t1_onnx # onnx_time.append(t1_onnxend) # average_time = np.sum(np.array(onnx_time[1:]))/999 # print("ONNX inference time:{}\n".format(average_time)) # print(y_onnx) #Prepare models for torch and tensorRT net = RFB(cfg=None, phase='test') trained_model = "./weights/RBF_Final.pth" device = torch.device("cuda") original_torch_time = [] tensorRT_time = [] model = load_model(net, trained_model, True).eval().to(device) x = torch.ones((1, 3, 224, 224)).cuda() size = [(480, 600), (720, 1280), (1920, 1080)] for sz in size: print("***** SIZE {} *****".format(sz)) x = torch.rand((1, 3, sz[0], sz[1])).to(device) x = x.to(device) model_trt = torch2trt(model, [x])