Exemple #1
0
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):
        ''' Create the mtcnn model '''
        pnet, rnet, onet = None, None, None

        if len(self.args.pnet_file) > 0:
            pnet = PNet(use_cuda=self.args.use_cuda)
            if self.args.use_cuda:
                pnet.load_state_dict(torch.load(self.args.pnet_file))
                pnet = torch.nn.DataParallel(
                    pnet, device_ids=self.args.gpu_ids).cuda()
            else:
                pnet.load_state_dict(torch.load(self.args.pnet_file,\
                                                map_location=lambda storage, loc: storage))
            pnet.eval()

        if len(self.args.rnet_file) > 0:
            rnet = RNet(use_cuda=self.args.use_cuda)
            if self.args.use_cuda:
                rnet.load_state_dict(torch.load(self.args.rnet_file))
                rnet = torch.nn.DataParallel(
                    rnet, device_ids=self.args.gpu_ids).cuda()
            else:
                rnet.load_state_dict(torch.load(self.args.rnet_file,\
                                                map_location=lambda storage, loc: storage))
            rnet.eval()

        if len(self.args.onet_file) > 0:
            onet = ONet(use_cuda=self.args.use_cuda)
            if self.args.use_cuda:
                onet.load_state_dict(torch.load(self.args.onet_file))
                onet = torch.nn.DataParallel(
                    onet, device_ids=self.args.gpu_ids).cuda()
            else:
                onet.load_state_dict(torch.load(self.args.onet_file, \
                                                map_location=lambda storage, loc: storage))
            onet.eval()

        self.pnet_detector = pnet
        self.rnet_detector = rnet
        self.onet_detector = onet
Exemple #3
0
def train_onet(model_store_path,
               end_epoch,
               imdb,
               batch_size,
               frequent=50,
               base_lr=0.01,
               lr_epoch_decay=[9],
               use_cuda=True,
               load=''):
    #create lr_list
    lr_epoch_decay.append(end_epoch + 1)
    lr_list = np.zeros(end_epoch)
    lr_t = base_lr
    for i in range(len(lr_epoch_decay)):
        if i == 0:
            lr_list[0:lr_epoch_decay[i] - 1] = lr_t
        else:
            lr_list[lr_epoch_decay[i - 1] - 1:lr_epoch_decay[i] - 1] = lr_t
        lr_t *= 0.1
    print(lr_list)

    if not os.path.exists(model_store_path):
        os.makedirs(model_store_path)

    lossfn = LossFn()
    net = ONet(is_train=True)
    if load != '':
        net.load_state_dict(torch.load(load))
        print('model loaded', load)
    net.train()
    print(use_cuda)
    if use_cuda:
        net.cuda()

    optimizer = torch.optim.Adam(net.parameters(), lr=base_lr)

    train_data = TrainImageReader(imdb, 48, batch_size, shuffle=True)

    for cur_epoch in range(1, end_epoch + 1):

        train_data.reset()
        for param in optimizer.param_groups:
            param['lr'] = lr_list[cur_epoch - 1]
        for batch_idx, (image, (gt_label, gt_bbox,
                                gt_landmark)) in enumerate(train_data):
            # print("batch id {0}".format(batch_idx))
            im_tensor = [
                image_tools.convert_image_to_tensor(image[i, :, :, :])
                for i in range(image.shape[0])
            ]
            im_tensor = torch.stack(im_tensor)

            im_tensor = Variable(im_tensor)
            gt_label = Variable(torch.from_numpy(gt_label).float())

            gt_bbox = Variable(torch.from_numpy(gt_bbox).float())
            gt_landmark = Variable(torch.from_numpy(gt_landmark).float())

            if use_cuda:
                im_tensor = im_tensor.cuda()
                gt_label = gt_label.cuda()
                gt_bbox = gt_bbox.cuda()
                gt_landmark = gt_landmark.cuda()

            cls_pred, box_offset_pred, landmark_offset_pred = net(im_tensor)

            # all_loss, cls_loss, offset_loss = lossfn.loss(gt_label=label_y,gt_offset=bbox_y, pred_label=cls_pred, pred_offset=box_offset_pred)

            cls_loss = lossfn.cls_loss(gt_label, cls_pred)
            box_offset_loss = lossfn.box_loss(gt_label, gt_bbox,
                                              box_offset_pred)
            #landmark_loss = lossfn.landmark_loss(gt_label,gt_landmark,landmark_offset_pred)

            all_loss = cls_loss * 0.8 + box_offset_loss * 0.6  #+landmark_loss*1.5

            if batch_idx % frequent == 0:
                accuracy = compute_accuracy(cls_pred, gt_label)

                show1 = accuracy.data.cpu().numpy()
                show2 = cls_loss.data.cpu().numpy()
                show3 = box_offset_loss.data.cpu().numpy()
                #show4 = landmark_loss.data.cpu().numpy()
                show5 = all_loss.data.cpu().numpy()

                #print("%s : Epoch: %d, Step: %d, accuracy: %s, det loss: %s, bbox loss: %s, landmark loss: %s, all_loss: %s, lr:%s "%(datetime.datetime.now(),cur_epoch,batch_idx, show1,show2,show3,show4,show5,base_lr))
                print(
                    "%s : Epoch: %d, Step: %d, accuracy: %s, det loss: %s, bbox loss: %s, all_loss: %s, lr:%s "
                    % (datetime.datetime.now(), cur_epoch, batch_idx, show1,
                       show2, show3, show5, lr_list[cur_epoch - 1]))

            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()

        torch.save(
            net.state_dict(),
            os.path.join(model_store_path, "onet_epoch_%d.pt" % cur_epoch))
        torch.save(
            net,
            os.path.join(model_store_path,
                         "onet_epoch_model_%d.pkl" % cur_epoch))