def create_mtcnn_net(p_model_path=None, r_model_path=None, o_model_path=None, use_cuda=True): pnet, rnet, onet = None, None, None if p_model_path is not None: pnet = PNet(use_cuda=use_cuda) pnet.load_state_dict(torch.load(p_model_path)) if (use_cuda): pnet.cuda() pnet.eval() if r_model_path is not None: rnet = RNet(use_cuda=use_cuda) rnet.load_state_dict(torch.load(r_model_path)) if (use_cuda): rnet.cuda() rnet.eval() if o_model_path is not None: onet = ONet(use_cuda=use_cuda) onet.load_state_dict(torch.load(o_model_path)) if (use_cuda): onet.cuda() onet.eval() return pnet, rnet, onet
def create_mtcnn_net(self, use_cuda=True): self.device = torch.device( "cuda" if use_cuda and torch.cuda.is_available() else "cpu") pnet = PNet() pnet.load_state_dict(model_zoo.load_url(model_urls['pnet'])) pnet.to(self.device).eval() onet = ONet() onet.load_state_dict(model_zoo.load_url(model_urls['onet'])) onet.to(self.device).eval() rnet = RNet() rnet.load_state_dict(model_zoo.load_url(model_urls['rnet'])) rnet.to(self.device).eval() return pnet, rnet, onet
def create_mtcnn_net(p_model_path=None, r_model_path=None, o_model_path=None, use_cuda=True): pnet, rnet, onet = None, None, None if p_model_path is not None: pnet = PNet(use_cuda=use_cuda) if (use_cuda): print('p_model_path:{0}'.format(p_model_path)) pnet.load_state_dict(torch.load(p_model_path)) pnet.cuda() else: # forcing all GPU tensors to be in CPU while loading pnet.load_state_dict( torch.load(p_model_path, map_location=lambda storage, loc: storage)) pnet.eval() if r_model_path is not None: rnet = RNet(use_cuda=use_cuda) if (use_cuda): print('r_model_path:{0}'.format(r_model_path)) rnet.load_state_dict(torch.load(r_model_path)) rnet.cuda() else: rnet.load_state_dict( torch.load(r_model_path, map_location=lambda storage, loc: storage)) rnet.eval() if o_model_path is not None: onet = ONet(use_cuda=use_cuda) if (use_cuda): print('o_model_path:{0}'.format(o_model_path)) onet.load_state_dict(torch.load(o_model_path)) onet.cuda() else: onet.load_state_dict( torch.load(o_model_path, map_location=lambda storage, loc: storage)) onet.eval() return pnet, rnet, onet
def create_mtcnn_net(self, use_cuda=True): self.device = torch.device( "cuda" if use_cuda and torch.cuda.is_available() else "cpu") pnet = PNet() #pnet.load_state_dict(torch.load(r'.\results\pnet\log_bs512_lr0.010_072402\check_point\model_050.pth')) pnet.load_state_dict(model_zoo.load_url(model_urls['pnet'])) pnet.to(self.device).eval() onet = ONet() #onet.load_state_dict(torch.load(r'.\results\onet\log_bs512_lr0.010_072602\check_point\model_050.pth')) onet.load_state_dict(model_zoo.load_url(model_urls['onet'])) onet.to(self.device).eval() rnet = RNet() #rnet.load_state_dict(torch.load(r'.\results\rnet\log_bs512_lr0.001_072502\check_point\model_050.pth')) rnet.load_state_dict(model_zoo.load_url(model_urls['rnet'])) rnet.to(self.device).eval() return pnet, rnet, onet