コード例 #1
0
def eval_net():
    # Assign imdb_name and imdbval_name according to args.dataset.
    if args.dataset == "voc":
        args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
        args.imdbval_name = "voc_2007_test"
    elif args.dataset == "coco":
        args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
        args.imdbval_name = "coco_2014_minival"
    # Import config
    if args.dataset == 'coco':
        cfg = (coco320, coco512)[args.input_size == 512]
    elif args.dataset == 'voc':
        cfg = (voc320, voc512)[args.input_size == 512]
    # Create imdb, roidb and blob_dataset
    print('Create or load an evaluted imdb.')
    imdb, roidb = combined_roidb(args.imdbval_name, False)
    imdb.competition_mode(on=True)
    print('{:d} roidb entries'.format(len(roidb)))

    blob_dataset = BlobDataset(imdb,
                               roidb,
                               transform=BaseTransform(cfg['min_dim'], MEANS),
                               target_normalization=True)

    # Construct networks.
    print('Construct {}_refinedet network.'.format(args.network))
    if args.network == 'vgg16':
        refinedet = VGGRefineDet(cfg['num_classes'], cfg)
    elif args.network == 'resnet101':
        refinedet = ResNetRefineDet(cfg['num_classes'], cfg)
    refinedet.create_architecture()
    # For CPU
    net = refinedet
    # For GPU/GPUs
    if args.cuda:
        if num_gpus > 1:
            net = torch.nn.DataParallel(refinedet)
        else:
            net = refinedet.cuda()
        cudnn.benchmark = True
    # Load weights
    net.load_weights(args.model_path)

    net.eval()
    print('Test RefineDet on:', args.imdbval_name)
    print('Using the specified args:')
    print(args)

    #data_loader = data.DataLoader(blob_dataset,
    #                              batch_size=1,
    #                              num_workers=0,
    #                              shuffle=False,
    #                              collate_fn=detection_collate,
    #                              pin_memory=True)
    # all detections are collected into:
    #    all_boxes[cls][image] = N x 5 array of detections in
    #    (x1, y1, x2, y2, score)
    num_images = len(imdb.image_index)
    # num_images = len(blob_dataset)
    num_classes = imdb.num_classes
    # num_object_classes + 1 ?
    print(num_classes)
    all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)]
    empty_array = np.transpose(np.array([[], [], [], [], []]), (1, 0))

    output_dir = get_output_dir(imdb, args.result_path)
    _t = {'im_detect': Timer(), 'misc': Timer()}
    det_file = os.path.join(output_dir, 'detections.pkl')

    # pdb.set_trace()
    for idx in range(num_images):
        img, gt, h, w = blob_dataset.pull_item(idx)
        input = Variable(img.unsqueeze(0), volatile=True)
        if args.cuda:
            input = input.cuda()

        # timers forward
        _t['im_detect'].tic()
        #start = time.time()
        # pdb.set_trace()
        detection = net(input)
        detect_time = _t['im_detect'].toc(average=True)
        #detect_time = time.time() - start
        print('im_detect: {:d}/{:d} {:.3f}s\n'.format(idx + 1, num_images,
                                                      detect_time))
        # skip jc = 0, because it's the background class
        for jc in range(1, num_classes):
            dets = detection[0, jc, :]
            mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
            dets = torch.masked_select(dets, mask).view(-1, 5)
            if dets.dim() > 0:
                boxes = dets[:, 1:]
                boxes[:, 0] *= w
                boxes[:, 2] *= w
                boxes[:, 1] *= h
                boxes[:, 3] *= h
                scores = dets[:, 0].cpu().numpy()
                cls_dets = np.hstack((boxes.cpu().numpy(),
                                      scores[:,
                                             np.newaxis])).astype(np.float32,
                                                                  copy=False)
                all_boxes[jc][idx] = cls_dets
            else:
                all_boxes[jc][idx] = empty_array

    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    print('Evaluating detections')
    imdb.evaluate_detections(all_boxes, output_dir)
コード例 #2
0
ファイル: train.py プロジェクト: 670373064/refinedet.pytorch
def main():
    images_path = './images.pkl'
    targets_path = './targets.pkl'
    
    args.dataset_root = COCO_ROOT
    cfg = coco320
    # data
    dataset = COCODetection(root=args.dataset_root,
                            transform=SSDAugmentation(cfg['min_dim'], MEANS))
    data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=False, collate_fn=detection_collate,
                                  pin_memory=True)
    ## load train data
    batch_iterator = iter(data_loader)
    # data_loader[0]
    images_list, targets_list = next(batch_iterator)
    pdb.set_trace()
    # images = torch.stack(images_list, 0)
    images_array_list = [Variable(cur_image).data.numpy() for cur_image in images_list]
    targets_array_list = [Variable(cur_targets).data.numpy() for cur_targets in targets_list]
    # store data
    fw_images = open(images_path, 'wb')
    cPickle.dump(images_array_list, fw_images)
    fw_images.close()
    fw_targets = open(targets_path, 'wb')
    cPickle.dump(targets_array_list, fw_targets)
    fw_targets.close()
    
    # load data
    images_array_list = cPickle.load(open(images_path, 'rb'))
    targets_array_list = cPickle.load(open(targets_path, 'rb'))
    
    images = torch.stack([torch.Tensor(cur_image) for cur_image in
                          images_array_list], 0)
    # targets = [torch.Tensor(cur_target) for cur_target in targets_array_list]
    targets = torch.stack([torch.Tensor(cur_target) for cur_target in
                           targets_array_list], 0)
    
    print(type(images))
    print(type(targets))
    
    refinedet = ResNetRefineDet(cfg['num_classes'], cfg)

    refinedet.create_architecture(
        os.path.join(args.save_folder, args.basenet), pretrained=True,
        fine_tuning=True)
    net = refinedet
    if args.cuda:
        net = torch.nn.DataParallel(refinedet).cuda()
        cudnn.benchmark = True

    params = net.state_dict()
    # for k, v in params.items():
    #     print(k)
    #     print(v.shape)
    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
                          weight_decay=args.weight_decay)
    net.train()
    print('Using the specified args:')
    print(args)

    if args.cuda:
        images = Variable(images.cuda())
        targets = Variable(targets.cuda())
        # targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
    else:
        images = Variable(images)
        targets = Variable(targets)

    # forward
    t0 = time.time()
    optimizer.zero_grad()
    bi_loss_loc, bi_loss_conf, multi_loss_loc, multi_loss_conf = \
        net(images, targets)
    loss = bi_loss_loc + bi_loss_conf + multi_loss_loc + multi_loss_conf
    loss.backward()
    optimizer.step()
    t1 = time.time()
    print('timer: %.4f sec.' % (t1 - t0))
コード例 #3
0
def train():
    # Assign imdb_name and imdbval_name according to args.dataset.
    if args.dataset == "voc":
        #args.imdb_name = "voc_2007_trainval"
        args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
        args.imdbval_name = "voc_2007_test"
    elif args.dataset == "coco":
        args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
        args.imdbval_name = "coco_2014_minival"
    # Import config
    if args.dataset == 'coco':
        cfg = (coco320, coco512)[args.input_size == 512]
    elif args.dataset == 'voc':
        cfg = (voc320, voc512)[args.input_size == 512]
    # Create imdb, roidb and blob_dataset
    print('Create or load an imdb.')
    imdb, roidb = combined_roidb(args.imdb_name)
    blob_dataset = BlobDataset(imdb,
                               roidb,
                               transform=SSDAugmentation(
                                   cfg['min_dim'], MEANS),
                               target_normalization=True)

    # Construct networks.
    print('Construct {}_refinedet network.'.format(args.network))
    if args.network == 'vgg16':
        refinedet = VGGRefineDet(cfg['num_classes'], cfg)
    elif args.network == 'resnet101':
        refinedet = ResNetRefineDet(cfg['num_classes'], cfg)
    refinedet.create_architecture(os.path.join(args.save_folder, args.basenet),
                                  pretrained=True,
                                  fine_tuning=True)
    #pdb.set_trace()
    # For CPU
    net = refinedet
    # For GPU/GPUs
    if args.cuda:
        if num_gpus > 1:
            net = torch.nn.DataParallel(refinedet)
        else:
            net = refinedet.cuda()
        cudnn.benchmark = True
    # Resume
    if args.resume_checkpoint:
        print('Resuming training, loading {}...'.format(
            args.resume_checkpoint))
        net.load_weights(args.resume_checkpoint)

    # pdb.set_trace()
    # params = net.state_dict()
    # for k, v in params.items():
    #     print(k)
    #     print(v.shape)
    optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
                          lr=args.lr,
                          momentum=args.momentum,
                          weight_decay=args.weight_decay)
    net.train()
    print('Training RefineDet on:', args.imdb_name)
    print('Using the specified args:')
    print(args)

    step_index = 0
    str_input_size = str(cfg['min_dim'])
    model_info = 'refinedet{0}_'.format(str_input_size) + args.dataset
    model_save_folder = os.path.join(args.save_folder, model_info)
    if not os.path.exists(model_save_folder):
        os.mkdir(model_save_folder)

    data_loader = data.DataLoader(blob_dataset,
                                  args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True,
                                  collate_fn=detection_collate,
                                  pin_memory=True)
    # Create batch iterator
    # Number of iterations in each epoch
    num_iter_per_epoch = len(blob_dataset) // args.batch_size
    # number of epoch, in case of resuming from args.start_iter
    num_epoch = (cfg['max_iter'] - args.start_iter) // num_iter_per_epoch
    iteration = args.start_iter
    arm_loss_loc = 0
    arm_loss_conf = 0
    odm_loss_loc = 0
    odm_loss_conf = 0
    for epoch in range(0, num_epoch):
        # pdb.set_trace()
        for i_batch, (images, targets) in enumerate(data_loader):
            if iteration in cfg['lr_steps']:
                step_index += 1
                adjust_learning_rate(optimizer, args.gamma, step_index)

            if args.cuda:
                images = Variable(images.cuda())
                targets = Variable(targets.cuda())
            else:
                images = Variable(images)
                targets = Variable(targets)
            # forward
            t0 = time.time()
            # backprop
            optimizer.zero_grad()
            bi_loss_loc, bi_loss_conf, multi_loss_loc, multi_loss_conf = \
                net(images, targets)
            loss = bi_loss_loc.mean() + bi_loss_conf.mean() + \
                   multi_loss_loc.mean() + multi_loss_conf.mean()
            loss.backward()
            optimizer.step()
            t1 = time.time()
            if num_gpus > 1:
                arm_loss_loc = bi_loss_loc.mean().data[0]
                arm_loss_conf = bi_loss_conf.mean().data[0]
                odm_loss_loc = multi_loss_loc.mean().data[0]
                odm_loss_conf = multi_loss_conf.mean().data[0]
            else:
                arm_loss_loc = bi_loss_loc.data[0]
                arm_loss_conf = bi_loss_conf.data[0]
                odm_loss_loc = multi_loss_loc.data[0]
                odm_loss_conf = multi_loss_conf.data[0]

            if iteration % 10 == 0:
                print('timer: %.4f sec.' % (t1 - t0))
                print('iter ' + repr(iteration) +
                      (' || ARM Loss Loc: %.4f  || ARM Loss Conf: %.4f' +
                       ' || ODM Loss Loc: %.4f  || ODM Loss Conf: %.4f' +
                       ' || Loss: %.4f ||') %
                      (arm_loss_loc, arm_loss_conf, odm_loss_loc,
                       odm_loss_conf, loss.data[0]) + ' ')


#                 print('iter ' + repr(iteration) +
#                       ' || Loss: %.4f ||' % (loss.data[0]) + ' ')
# print('iter ' + repr(iteration) +
#       ' || Loss: %.4f ||' % (loss.data[0]), end=' ')

            if iteration != 0 and iteration % 10000 == 0:
                #if iteration != 0 and iteration % cfg['checkpoint_step'] == 0:
                print('Saving state, iter:', iteration)
                torch.save(
                    refinedet.state_dict(),
                    os.path.join(model_save_folder,
                                 model_info + '_' + repr(iteration) + '.pth'))

            iteration += 1

    torch.save(refinedet.state_dict(),
               os.path.join(args.save_folder, args.dataset + '.pth'))
コード例 #4
0
def train():
    # Assign imdb_name and imdbval_name according to args.dataset.
    if args.dataset == "voc":
        # args.imdb_name = "voc_2007_trainval"
        args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
        args.imdbval_name = "voc_2007_test"
    elif args.dataset == "coco":
        args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
        args.imdbval_name = "coco_2014_minival"
    # Import config
    if args.dataset == 'coco':
        cfg = (coco320, coco512)[args.input_size == 512]
    elif args.dataset == 'voc':
        cfg = (voc320, voc512)[args.input_size == 512]
    # Create imdb, roidb and blob_dataset
    print('Create or load an imdb.')
    imdb, roidb = combined_roidb(args.imdb_name)
    blob_dataset = BlobDataset(imdb,
                               roidb,
                               transform=SSDAugmentation(
                                   cfg['min_dim'], MEANS),
                               target_normalization=True)

    # Construct networks.
    print('Construct {}_refinedet network.'.format(args.network))
    if args.network == 'vgg16':
        refinedet = VGGRefineDet(cfg['num_classes'], cfg)
    elif args.network == 'resnet101':
        refinedet = ResNetRefineDet(cfg['num_classes'], cfg)

    pretrained_model = os.path.join(path_cfg.DATA_DIR, args.pretrained_folder,
                                    args.base_model)
    refinedet.create_architecture(pretrained_model,
                                  pretrained=True,
                                  fine_tuning=True)
    # For CPU
    net = refinedet
    # For GPU/GPUs
    if args.cuda:
        net = refinedet.cuda()
        if num_gpus > 1:
            net = torch.nn.DataParallel(net)
        cudnn.benchmark = True
    # Resume
    if args.resume_checkpoint:
        print('Resuming training, loading {}...'.format(
            args.resume_checkpoint))
        net.load_weights(args.resume_checkpoint)

    # pdb.set_trace()
    # params = net.state_dict()
    # for k, v in params.items():
    #     print(k)
    #     print(v.shape)
    optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
                          lr=args.lr,
                          momentum=args.momentum,
                          weight_decay=args.weight_decay)
    net.train()
    print('Training RefineDet on:', args.imdb_name)
    print('Using the specified args:')
    print(args)

    output_folder = os.path.join(path_cfg.OUTPUT_DIR, args.output_folder)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    str_input_size = str(cfg['min_dim'])
    model_info = 'refinedet{}_{}'.format(str_input_size, args.dataset)
    model_output_folder = os.path.join(output_folder,
                                       '{}'.format(args.network), model_info)
    if not os.path.exists(model_output_folder):
        os.makedirs(model_output_folder)

    data_loader = data.DataLoader(blob_dataset,
                                  args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True,
                                  collate_fn=detection_collate,
                                  pin_memory=True)
    # Create batch iterator
    # Number of iterations in each epoch
    base_batch_size = 32
    batch_multiplier = float(args.batch_size) / base_batch_size
    # number of epoch, in case of resuming from args.start_iter
    # fixed number of epoches, no matter what args.batch_size is
    num_epoch = (cfg['max_iter'] - args.start_iter) // (len(blob_dataset) //
                                                        base_batch_size)
    base_iteration = args.start_iter
    lr_steps = [int(x / batch_multiplier) for x in cfg['lr_steps']]
    actual_iteration = 0
    decay_step = 0
    # print('num_epoch: {}, batch_multiplier: {}, maximum base_iteration {} actual_iteration: {}')
    for epoch in range(0, num_epoch):
        # pdb.set_trace()
        t0 = time.time()
        for i_batch, (images, targets) in enumerate(data_loader):
            if actual_iteration in lr_steps:
                decay_step += 1
                adjust_learning_rate(optimizer, args.gamma, decay_step)

            if args.cuda:
                images = Variable(images.cuda())
                targets = Variable(targets.cuda())
            else:
                images = Variable(images)
                targets = Variable(targets)
            t1_data = time.time()
            # forward and backprop
            optimizer.zero_grad()
            bi_loss_loc, bi_loss_conf, multi_loss_loc, multi_loss_conf = \
                net(images, targets)
            loss = bi_loss_loc.mean() + bi_loss_conf.mean() + \
                   multi_loss_loc.mean() + multi_loss_conf.mean()
            # loss = bi_loss_loc.mean() + bi_loss_conf.mean()
            loss.backward()
            optimizer.step()
            t1 = time.time()
            if num_gpus > 1:
                arm_loss_loc = bi_loss_loc.mean().item()
                arm_loss_conf = bi_loss_conf.mean().item()
                odm_loss_loc = multi_loss_loc.mean().item()
                odm_loss_conf = multi_loss_conf.mean().item()
            else:
                arm_loss_loc = bi_loss_loc.item()
                arm_loss_conf = bi_loss_conf.item()
                odm_loss_loc = multi_loss_loc.item()
                odm_loss_conf = multi_loss_conf.item()

            if actual_iteration % 10 == 0:
                print('timer: %.4f sec, data loading timer: %.4f sec' %
                      (t1 - t0, t1_data - t0))
                print('iter ' + repr(actual_iteration) +
                      (' || ARM Loss Loc: %.4f  || ARM Loss Conf: %.4f' +
                       ' || ODM Loss Loc: %.4f  || ODM Loss Conf: %.4f' +
                       ' || Loss: %.4f ||') %
                      (arm_loss_loc, arm_loss_conf, odm_loss_loc,
                       odm_loss_conf, loss.item()) + ' ')
            # save checkpoint.
            if actual_iteration != 0 and actual_iteration % (int(
                    10000 / batch_multiplier)) == 0:
                print('Saving state, iter:', base_iteration)
                torch.save(
                    refinedet.state_dict(),
                    os.path.join(
                        model_output_folder, '_'.join([
                            args.network, model_info,
                            repr(base_iteration) + '.pth'
                        ])))
            # update counts.
            actual_iteration += 1
            base_iteration = int(actual_iteration * batch_multiplier)
            t0 = time.time()

    torch.save(
        refinedet.state_dict(),
        os.path.join(model_output_folder,
                     '_'.join([args.network, model_info + '.pth'])))