Esempio n. 1
0
def create():
    print("==========\nArgs:{}\n==========".format(args))
    print("Goal: randomly split data for {} times, {:.1%} for training and the rest for testing".format(args.num_splits, args.train_percent))
    print("Loading dataset from {}".format(args.dataset))
    dataset = h5py.File(args.dataset, 'r')
    keys = dataset.keys()
    num_videos = len(keys)
    num_train = int(math.ceil(num_videos * args.train_percent))
    num_test = num_videos - num_train
    print("Split breakdown: # total videos {}. # train videos {}. # test videos {}".format(num_videos, num_train, num_test))
    splits = []

    for split_idx in range(args.num_splits):
        train_keys, test_keys = split_random(keys, num_videos, num_train)
        splits.append({
            'train_keys': train_keys,
            'test_keys': test_keys,
            })

    if args.save_name is None:
        save_name = dir_utils.get_stem(args.dataset)
    else:
        save_name = args.save_name
    save_dir = dir_utils.get_dir(args.save_dir)
    saveto = osp.join(save_dir, save_name + '.json')
    write_json(splits, saveto)
    print("Splits saved to {}".format(saveto))

    dataset.close()
Esempio n. 2
0
def parse_c3dfiles(path):
    """
    Parse directories holding extracted frames from standard benchmarks
    """
    print('parse frames under folder {}'.format(path))
    feature_files = glob.glob(os.path.join(path, '*.mat'))

    # check RGB
    frame_dict = {}
    for i, s_feature_file in enumerate(feature_files):
        s_feature = scipy.io.loadmat(s_feature_file)['relu6']
        n_frames = s_feature.shape[0]
        k = dir_utils.get_stem(s_feature_file)

        frame_dict[k] = (s_feature_file, n_frames)

    print('frame folder analysis done')
    return frame_dict
def config_model(char_to_id, tag_to_id, args=None):
    config = OrderedDict()
    config["model_type"] = dir_utils.get_stem(__file__)
    config["num_chars"] = len(char_to_id)
    config["char_dim"] = args.char_dim
    config["num_tags"] = len(tag_to_id)
    config["seg_dim"] = args.seg_dim
    config["lstm_dim"] = args.hidden_dim
    config["batch_size"] = args.batch_size

    config["emb_file"] = 'data/vec.txt'
    # config["clip"] = FLAGS.clip
    config["dropout"] = args.dropout  # Update: 0 means not drop out at all
    config["tag_schema"] = args.tag_schema

    config['train_file'] = 'data/example.train'
    config['dev_file'] = 'data/example.dev'
    config['test_file'] = 'data/example.test'
    config['map_file'] = 'map.pkl'

    return config
Esempio n. 4
0
def main():
    PCA_info = pkl.load(open('pca_c3dd_fc7_val_annot_thumos14.pkl', 'rb'))
    x_mean = PCA_info['x_mean']
    U = PCA_info['U']
    num_red_dim = 500

    feature_directory = '/home/zwei/datasets/THUMOS14/features/c3dd-fc7'
    feature_files = glob.glob(os.path.join(feature_directory, '*.mat'))
    target_directory = dir_utils.get_dir(
        '/home/zwei/datasets/THUMOS14/features/c3dd-fc7-red500')

    print time.ctime(), 'start'
    pbar = progressbar.ProgressBar(max_value=len(feature_files))
    for feat_idx, s_feature_file in enumerate(feature_files):
        s_file_stem = dir_utils.get_stem(s_feature_file)
        output_file = os.path.join(target_directory,
                                   '{:s}.npy'.format(s_file_stem))
        pbar.update(feat_idx)
        s_feature = scipy.io.loadmat(s_feature_file)['fc7']
        s_feature_red = np.dot(s_feature - x_mean, U[:, :num_red_dim])
        np.save(output_file, s_feature_red)
Esempio n. 5
0
def parse_c3ddfiles(path):
    """
    Parse directories holding extracted frames from standard benchmarks
    """
    print('parse frames under folder {}'.format(path))
    feature_files = glob.glob(os.path.join(path, '*.npy'))

    frame_directories = glob.glob(
        os.path.join('/home/zwei/datasets/THUMOS14/frame', '*/'))
    frame_directory = '/home/zwei/datasets/THUMOS14/frame'
    # check RGB
    frame_dict = {}
    for i, s_feature_file in enumerate(feature_files):
        s_feature = np.load(s_feature_file)
        n_frames = s_feature.shape[0]

        k = dir_utils.get_stem(s_feature_file)
        s_frames = len(glob.glob(os.path.join(frame_directory, k, '*.jpg')))
        assert s_frames == n_frames, 'BUG'
        frame_dict[k] = (s_feature_file, n_frames)

    print('frame folder analysis done')
    return frame_dict
Esempio n. 6
0
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = dir_utils.get_dir(
        os.path.join(
            project_root, 'ckpts',
            '{:s}-{:s}-{:s}-split-{:d}-claweight-{:s}-{:.1f}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-seqlen{:d}-samplerate-{:d}-{:s}-{:s}'
            .format(script_name_stem, args.dataset, args.eval_metrics,
                    args.split, str(args.set_cls_weight), args.cls_pos_weight,
                    args.hassign_thres, args.alpha, args.hidden_dim,
                    args.dropout, args.seq_len, args.sample_rate,
                    loss_type[args.EMD], match_type[args.hmatch])))
    log_file = os.path.join(save_directory,
                            'log-{:s}.txt'.format(dir_utils.get_date_str()))
    logger = log_utils.get_logger(log_file)
    log_utils.print_config(vars(args), logger)

    model = PointerNetwork(input_dim=args.input_dim,
                           embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim,
                           max_decoding_len=args.net_outputs,
                           dropout=args.dropout,
                           n_enc_layers=2,
                           output_classes=2)
    hassign_thres = args.hassign_thres
    logger.info("Number of Params\t{:d}".format(
        sum([p.data.nelement() for p in model.parameters()])))
    logger.info('Saving logs to {:s}'.format(log_file))

    if args.resume is not None:

        ckpt_idx = 48

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(
            ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename,
                                map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        logger.info("=> loading checkpoint '{}', current iou: {:.04f}".format(
            ckpt_filename, train_iou))

    model = cuda_model.convertModel2Cuda(model,
                                         gpu_id=args.gpu_id,
                                         multiGpu=args.multiGpu)
    # get train/val split
    if args.dataset == 'SumMe':
        train_val_test_perms = np.arange(25)
    elif args.dataset == 'TVSum':
        train_val_test_perms = np.arange(50)
    # fixed permutation
    random.Random(0).shuffle(train_val_test_perms)
    train_val_test_perms = train_val_test_perms.reshape([5, -1])
    train_val_perms = np.delete(train_val_test_perms, args.split,
                                0).reshape([-1])
    train_perms = train_val_perms[:17]
    val_perms = train_val_perms[17:]
    test_perms = train_val_test_perms[args.split]
    logger.info(" training split: " + str(train_perms))
    logger.info(" val split: " + str(val_perms))
    logger.info(" test split: " + str(test_perms))

    if args.location == 'home':
        data_path = os.path.join(os.path.expanduser('~'), 'datasets')
    else:
        data_path = os.path.join('/nfs/%s/boyu/SDN' % (args.location),
                                 'datasets')
    train_dataset = vsSumLoader3_c3dd.cDataset(dataset_name=args.dataset,
                                               split='train',
                                               seq_length=args.seq_len,
                                               overlap=0.9,
                                               sample_rate=[args.sample_rate],
                                               train_val_perms=train_perms,
                                               data_path=data_path)
    val_evaluator = Evaluator.Evaluator(dataset_name=args.dataset,
                                        split='val',
                                        seq_length=args.seq_len,
                                        overlap=0.9,
                                        sample_rate=[args.sample_rate],
                                        sum_budget=0.15,
                                        train_val_perms=val_perms,
                                        eval_metrics=args.eval_metrics,
                                        data_path=data_path)
    test_evaluator = Evaluator.Evaluator(dataset_name=args.dataset,
                                         split='test',
                                         seq_length=args.seq_len,
                                         overlap=0.9,
                                         sample_rate=[args.sample_rate],
                                         sum_budget=0.15,
                                         train_val_perms=test_perms,
                                         eval_metrics=args.eval_metrics,
                                         data_path=data_path)

    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    # val_dataloader = DataLoader(val_dataset,
    #                               batch_size=args.batch_size,
    #                               shuffle=False,
    #                               num_workers=4)

    model_optim = optim.Adam(filter(lambda p: p.requires_grad,
                                    model.parameters()),
                             lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim,
                                                           'min',
                                                           patience=10)

    alpha = args.alpha
    # cls_weights = torch.FloatTensor([0.2, 1.0]).cuda()
    if args.set_cls_weight:
        cls_weights = torch.FloatTensor([
            1. * train_dataset.n_positive_train_samples /
            train_dataset.n_total_train_samples, args.cls_pos_weight
        ]).cuda()
    else:
        cls_weights = torch.FloatTensor([0.5, 0.5]).cuda()
    logger.info(" total: {:d}, total pos: {:d}".format(
        train_dataset.n_total_train_samples,
        train_dataset.n_positive_train_samples))
    logger.info(" classify weight: " + str(cls_weights[0]) +
                str(cls_weights[1]))
    for epoch in range(args.start_epoch, args.nof_epoch + args.start_epoch):
        total_losses = AverageMeter()
        loc_losses = AverageMeter()
        cls_losses = AverageMeter()
        Accuracy = AverageMeter()
        IOU = AverageMeter()
        ordered_IOU = AverageMeter()
        model.train()
        pbar = progressbar.ProgressBar(max_value=len(train_dataloader))
        for i_batch, sample_batch in enumerate(train_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            start_indices = Variable(sample_batch[1])
            end_indices = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])
            # seq_labels = Variable(sample_batch[3])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                start_indices = start_indices.cuda()
                end_indices = end_indices.cuda()

            gt_positions = torch.stack([start_indices, end_indices], dim=-1)

            head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                feature_batch)

            pred_positions = torch.stack([head_positions, tail_positions],
                                         dim=-1)
            if args.hmatch:
                assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_v2(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)

            else:
                assigned_scores, assigned_locations = f_match.Assign_Batch(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)
                _, _, total_valid, total_iou = h_match.Assign_Batch_v2(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)

            if total_valid > 0:
                IOU.update(total_iou / total_valid, total_valid)

            assigned_scores = Variable(torch.LongTensor(assigned_scores),
                                       requires_grad=False)
            assigned_locations = Variable(torch.LongTensor(assigned_locations),
                                          requires_grad=False)
            if use_cuda:
                assigned_scores = assigned_scores.cuda()
                assigned_locations = assigned_locations.cuda()

            cls_scores = cls_scores.contiguous().view(-1,
                                                      cls_scores.size()[-1])
            assigned_scores = assigned_scores.contiguous().view(-1)

            cls_loss = F.cross_entropy(cls_scores,
                                       assigned_scores,
                                       weight=cls_weights)

            if total_valid > 0:
                assigned_head_positions = assigned_locations[:, :, 0]
                assigned_head_positions = assigned_head_positions.contiguous(
                ).view(-1)
                #
                assigned_tail_positions = assigned_locations[:, :, 1]
                assigned_tail_positions = assigned_tail_positions.contiguous(
                ).view(-1)

                head_pointer_probs = head_pointer_probs.contiguous().view(
                    -1,
                    head_pointer_probs.size()[-1])
                tail_pointer_probs = tail_pointer_probs.contiguous().view(
                    -1,
                    tail_pointer_probs.size()[-1])

                assigned_head_positions = torch.masked_select(
                    assigned_head_positions, assigned_scores.byte())
                assigned_tail_positions = torch.masked_select(
                    assigned_tail_positions, assigned_scores.byte())

                head_pointer_probs = torch.index_select(
                    head_pointer_probs,
                    dim=0,
                    index=assigned_scores.nonzero().squeeze(1))
                tail_pointer_probs = torch.index_select(
                    tail_pointer_probs,
                    dim=0,
                    index=assigned_scores.nonzero().squeeze(1))

                if args.EMD:
                    assigned_head_positions = to_one_hot(
                        assigned_head_positions, args.seq_len)
                    assigned_tail_positions = to_one_hot(
                        assigned_tail_positions, args.seq_len)

                    prediction_head_loss = EMD_L2(head_pointer_probs,
                                                  assigned_head_positions,
                                                  needSoftMax=True)
                    prediction_tail_loss = EMD_L2(tail_pointer_probs,
                                                  assigned_tail_positions,
                                                  needSoftMax=True)
                else:
                    prediction_head_loss = F.cross_entropy(
                        head_pointer_probs, assigned_head_positions)
                    prediction_tail_loss = F.cross_entropy(
                        tail_pointer_probs, assigned_tail_positions)
                loc_losses.update(
                    prediction_head_loss.data.item() +
                    prediction_tail_loss.data.item(), feature_batch.size(0))
                total_loss = alpha * (prediction_head_loss +
                                      prediction_tail_loss) + cls_loss
            else:
                total_loss = cls_loss

            model_optim.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
            model_optim.step()
            cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
            total_losses.update(total_loss.item(), feature_batch.size(0))

        logger.info(
            "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tcls-Accuracy:{:.4f}\tloc-Avg-IOU:{:.4f}\t topIOU:{:.4f}"
            .format(epoch, model_optim.param_groups[0]['lr'], total_losses.avg,
                    cls_losses.avg, loc_losses.avg, Accuracy.avg, IOU.avg,
                    ordered_IOU.avg))

        optim_scheduler.step(total_losses.avg)

        model.eval()

        # IOU = AverageMeter()
        # pbar = progressbar.ProgressBar(max_value=len(val_evaluator))
        # for i_batch, sample_batch in enumerate(val_dataloader):
        #     pbar.update(i_batch)

        #     feature_batch = Variable(sample_batch[0])
        #     start_indices = Variable(sample_batch[1])
        #     end_indices = Variable(sample_batch[2])
        #     gt_valids = Variable(sample_batch[3])
        #     # valid_indices = Variable(sample_batch[3])

        #     if use_cuda:
        #         feature_batch = feature_batch.cuda()
        #         start_indices = start_indices.cuda()
        #         end_indices = end_indices.cuda()

        #     gt_positions = torch.stack([start_indices, end_indices], dim=-1)

        #     head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
        #         feature_batch)#Update: compared to the previous version, we now update the matching rules

        #     pred_positions = torch.stack([head_positions, tail_positions], dim=-1)
        #     pred_scores = cls_scores[:, :, -1]
        #     #TODO: should NOT change here for evaluation!
        #     assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_v2(gt_positions, pred_positions, gt_valids, thres=hassign_thres)
        #     if total_valid>0:
        #         IOU.update(total_iou / total_valid, total_valid)

        val_F1s = val_evaluator.Evaluate(model)
        test_F1s = test_evaluator.Evaluate(model)

        logger.info("Val -- Epoch :{:06d}, LR: {:.6f},\tF1s:{:.4f}".format(
            epoch, model_optim.param_groups[0]['lr'], val_F1s))
        logger.info("Test -- Epoch :{:06d},\tF1s:{:.4f}".format(
            epoch, test_F1s))

        if epoch % 1 == 0:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'loss': total_losses.avg,
                    'cls_loss': cls_losses.avg,
                    'loc_loss': loc_losses.avg,
                    'IoU': IOU.avg,
                    'val_F1s': val_F1s,
                    'test_F1s': test_F1s
                }, (epoch + 1),
                file_direcotry=save_directory)
Esempio n. 7
0
import scipy.io as sio
import sys, os
from PyUtils import dir_utils
import numpy as np
import glob
import progressbar

feature_directory = '/home/zwei/datasets/THUMOS14/features/c3d_feat_dense'
target_directory = dir_utils.get_dir(
    '/home/zwei/datasets/THUMOS14/features/c3d-dense')

feature_lists = glob.glob(os.path.join(feature_directory, '*.mat'))

# video_name = 'video_validation_0000940.mat'
pbar = progressbar.ProgressBar(max_value=len(feature_lists))
for file_id, s_feature_path in enumerate(feature_lists):
    pbar.update(file_id)
    # s_feature_name = os.path.basename(s_feature_path)
    s_feature_stem = dir_utils.get_stem(s_feature_path)
    # feature_path = os.path.join(, video_name)
    s_full_feature = sio.loadmat(s_feature_path)['relu6']
    np.save(os.path.join(target_directory, '{:s}.npy'.format(s_feature_stem)),
            s_full_feature)
print("DEBUG")
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)

    # Pretty print the run args
    pp.pprint(vars(args))

    model = PointerNetwork(input_dim=args.input_dim, embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim, max_decoding_len=args.net_outputs, dropout=args.dropout, n_enc_layers=2)
    hassign_thres = args.hassign_thres
    print("Number of Params\t{:d}".format(sum([p.data.nelement() for p in model.parameters()])))
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = os.path.join(project_root, 'ckpts', '{:s}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-seqlen{:d}-ckpt'.
                                  format(script_name_stem, hassign_thres, args.alpha, args.hidden_dim, args.dropout, args.seq_len))
    print("Save ckpt to {:s}".format(save_directory))

    if args.resume is not None:

        ckpt_idx = 3

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename, map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        print("=> loading checkpoint '{}', current iou: {:.04f}".format(ckpt_filename, train_iou))


    model = cuda_model.convertModel2Cuda(model, gpu_id=args.gpu_id, multiGpu=args.multiGpu)

    train_dataset  = cDataset(seq_length=args.seq_len, overlap=0.9, sample_rate=[4], dataset_split='train', rdDrop=True, rdOffset=True)

    val_dataset = cDataset(seq_length=args.seq_len, overlap=0.9, sample_rate=[4], dataset_split='val', rdDrop=False, rdOffset=False)


    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=False,
                                  num_workers=4)

    model_optim = optim.Adam(filter(lambda p:p.requires_grad,  model.parameters()), lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim, 'min', patience=20)

    alpha=args.alpha
    cls_weights = torch.FloatTensor([0.05, 1.0]).cuda()
    for epoch in range(args.start_epoch, args.nof_epoch+args.start_epoch):
            total_losses = AverageMeter()
            loc_losses = AverageMeter()
            cls_losses = AverageMeter()
            Accuracy = AverageMeter()
            IOU = AverageMeter()
            ordered_IOU = AverageMeter()
            model.train()
            pbar = progressbar.ProgressBar(max_value=len(train_dataloader))
            for i_batch, sample_batch in enumerate(train_dataloader):
                pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                start_indices = Variable(sample_batch[1])
                end_indices = Variable(sample_batch[2])
                gt_valids = Variable(sample_batch[3])

                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    start_indices = start_indices.cuda()
                    end_indices = end_indices.cuda()

                gt_positions = torch.stack([start_indices, end_indices], dim=-1)

                head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(feature_batch)

                pred_positions = torch.stack([head_positions, tail_positions], dim=-1)

                assigned_scores, assigned_locations = h_assign.Assign_Batch(gt_positions, pred_positions, gt_valids, thres=hassign_thres)
                # if np.sum(assigned_scores) > 1:
                #     print("DEBUG")
                # correct_predictions = np.sum(assigned_scores[:,:args.n_outputs])
                # cls_rate = correct_predictions*1./np.sum(assigned_scores)
                if np.sum(assigned_scores)>=1:
                    iou_rate, effective_positives = Metrics.get_avg_iou2(np.reshape(pred_positions.data.cpu().numpy(), (-1, 2)),
                                                   np.reshape(assigned_locations, (-1, 2)), np.reshape(assigned_scores,
                                                                                                       assigned_scores.shape[
                                                                                                           0] *
                                                                                                       assigned_scores.shape[
                                                                                                           1]))

                    IOU.update(iou_rate/(effective_positives), effective_positives)
                    # ordered_IOU.update(ordered_iou_rate/(args.batch_size*args.n_outputs),args.batch_size*args.n_outputs)

                    # n_effective_batches += 1

                assigned_scores = Variable(torch.LongTensor(assigned_scores),requires_grad=False)
                assigned_locations = Variable(torch.LongTensor(assigned_locations), requires_grad=False)
                if use_cuda:
                    assigned_scores = assigned_scores.cuda()
                    assigned_locations = assigned_locations.cuda()

                cls_scores = cls_scores.contiguous().view(-1, cls_scores.size()[-1])
                assigned_scores = assigned_scores.contiguous().view(-1)


                cls_loss = F.cross_entropy(cls_scores, assigned_scores, weight=cls_weights)

                if torch.sum(assigned_scores)>0:
                    # print("HAHA")
                    assigned_head_positions = assigned_locations[:,:,0]
                    assigned_head_positions = assigned_head_positions.contiguous().view(-1)
                    #
                    assigned_tail_positions = assigned_locations[:,:,1]
                    assigned_tail_positions = assigned_tail_positions.contiguous().view(-1)


                    head_pointer_probs = head_pointer_probs.contiguous().view(-1, head_pointer_probs.size()[-1])
                    tail_pointer_probs = tail_pointer_probs.contiguous().view(-1, tail_pointer_probs.size()[-1])


                    # mask here: if there is non in assigned scores, no need to compute ...

                    assigned_head_positions = torch.masked_select(assigned_head_positions, assigned_scores.byte())
                    assigned_tail_positions = torch.masked_select(assigned_tail_positions, assigned_scores.byte())

                    head_pointer_probs = torch.index_select(head_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))
                    tail_pointer_probs = torch.index_select(tail_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))

                    assigned_head_positions = to_one_hot(assigned_head_positions, args.seq_len)
                    assigned_tail_positions = to_one_hot(assigned_tail_positions, args.seq_len)

                    prediction_head_loss = EMD_L2(head_pointer_probs, assigned_head_positions, needSoftMax=True)
                    prediction_tail_loss = EMD_L2(tail_pointer_probs, assigned_tail_positions, needSoftMax=True)
                    loc_losses.update(prediction_head_loss.data.item() + prediction_tail_loss.data.item(),
                                      feature_batch.size(0))
                    total_loss = alpha * (prediction_head_loss + prediction_tail_loss) + cls_loss
                else:
                    total_loss = cls_loss

                model_optim.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
                model_optim.step()
                cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
                total_losses.update(total_loss.item(), feature_batch.size(0))


            print(
                "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tcls-Accuracy:{:.4f}\tloc-Avg-IOU:{:.4f}\t topIOU:{:.4f}".format(
                    epoch,
                    model_optim.param_groups[0]['lr'], total_losses.avg, cls_losses.avg, loc_losses.avg, Accuracy.avg, IOU.avg, ordered_IOU.avg))

            optim_scheduler.step(total_losses.avg)

            model.eval()
            total_losses = AverageMeter()
            loc_losses = AverageMeter()
            cls_losses = AverageMeter()
            Accuracy = AverageMeter()
            IOU = AverageMeter()
            ordered_IOU = AverageMeter()
            pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
            for i_batch, sample_batch in enumerate(val_dataloader):
                pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                start_indices = Variable(sample_batch[1])
                end_indices = Variable(sample_batch[2])
                gt_valids = Variable(sample_batch[3])
                # valid_indices = Variable(sample_batch[3])

                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    start_indices = start_indices.cuda()
                    end_indices = end_indices.cuda()

                gt_positions = torch.stack([start_indices, end_indices], dim=-1)

                head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                    feature_batch)

                pred_positions = torch.stack([head_positions, tail_positions], dim=-1)

                assigned_scores, assigned_locations = h_assign.Assign_Batch(gt_positions, pred_positions, gt_valids, thres=hassign_thres)
                # if np.sum(assigned_scores) > 1:
                #     print("DEBUG")
                # correct_predictions = np.sum(assigned_scores[:,:args.n_outputs])
                # cls_rate = correct_predictions*1./np.sum(assigned_scores)
                if np.sum(assigned_scores) >= 1:
                    iou_rate, effective_positives = Metrics.get_avg_iou2(
                        np.reshape(pred_positions.data.cpu().numpy(), (-1, 2)),
                        np.reshape(assigned_locations, (-1, 2)), np.reshape(assigned_scores,
                                                                            assigned_scores.shape[
                                                                                0] *
                                                                            assigned_scores.shape[
                                                                                1]))
                    IOU.update(iou_rate / (effective_positives), effective_positives)


                assigned_scores = Variable(torch.LongTensor(assigned_scores), requires_grad=False)
                assigned_locations = Variable(torch.LongTensor(assigned_locations), requires_grad=False)
                if use_cuda:
                    assigned_scores = assigned_scores.cuda()
                    assigned_locations = assigned_locations.cuda()

                cls_scores = cls_scores.contiguous().view(-1, cls_scores.size()[-1])
                assigned_scores = assigned_scores.contiguous().view(-1)

                cls_loss = F.cross_entropy(cls_scores, assigned_scores, weight=cls_weights)

                if torch.sum(assigned_scores)>0:
                    # print("HAHA")
                    assigned_head_positions = assigned_locations[:,:,0]
                    assigned_head_positions = assigned_head_positions.contiguous().view(-1)
                    #
                    assigned_tail_positions = assigned_locations[:,:,1]
                    assigned_tail_positions = assigned_tail_positions.contiguous().view(-1)


                    head_pointer_probs = head_pointer_probs.contiguous().view(-1, head_pointer_probs.size()[-1])
                    tail_pointer_probs = tail_pointer_probs.contiguous().view(-1, tail_pointer_probs.size()[-1])


                    # mask here: if there is non in assigned scores, no need to compute ...

                    assigned_head_positions = torch.masked_select(assigned_head_positions, assigned_scores.byte())
                    assigned_tail_positions = torch.masked_select(assigned_tail_positions, assigned_scores.byte())

                    head_pointer_probs = torch.index_select(head_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))
                    tail_pointer_probs = torch.index_select(tail_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))

                    assigned_head_positions = to_one_hot(assigned_head_positions, args.seq_len)
                    assigned_tail_positions = to_one_hot(assigned_tail_positions, args.seq_len)

                    prediction_head_loss = EMD_L2(head_pointer_probs, assigned_head_positions, needSoftMax=True)
                    prediction_tail_loss = EMD_L2(tail_pointer_probs, assigned_tail_positions, needSoftMax=True)
                    loc_losses.update(prediction_head_loss.data.item() + prediction_tail_loss.data.item(),
                                      feature_batch.size(0))
                    total_loss = alpha * (prediction_head_loss + prediction_tail_loss) + cls_loss
                else:
                    total_loss = cls_loss

                cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
                total_losses.update(total_loss.item(), feature_batch.size(0))

            print(
                "Val -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tcls-Accuracy:{:.4f}\tloc-Avg-IOU:{:.4f}\t topIOU:{:.4f}".format(
                    epoch,
                    model_optim.param_groups[0]['lr'], total_losses.avg, cls_losses.avg, loc_losses.avg, Accuracy.avg,
                    IOU.avg, ordered_IOU.avg))


            if epoch % 1 == 0:
                save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'loss':total_losses.avg,
            'cls_loss': cls_losses.avg,
            'loc_loss': loc_losses.avg,
            'IoU': IOU.avg}, (epoch+1), file_direcotry=save_directory)
Esempio n. 9
0
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)
    script_name_stem = dir_utils.get_stem(__file__)
    if args.resume is None:

        save_directory = dir_utils.get_dir(os.path.join(project_root, 'ckpts', '{:s}'.format(args.dataset), '{:s}-{:s}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-seqlen{:d}-{:s}-{:s}'.
                                      format(script_name_stem, args.sufix, args.hassign_thres, args.alpha, args.hidden_dim, args.dropout, args.seq_len, 'L2', match_type[args.hmatch])))
    else:
        save_directory = args.resume

    log_file = os.path.join(save_directory, 'log-{:s}.txt'.format(dir_utils.get_date_str()))
    logger = log_utils.get_logger(log_file)
    log_utils.print_config(vars(args), logger)

    model = PointerNetwork(input_dim=args.input_dim, embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim, max_decoding_len=args.net_outputs, dropout=args.dropout, n_enc_layers=2, output_classes=2)
    logger.info("Number of Params\t{:d}".format(sum([p.data.nelement() for p in model.parameters()])))
    logger.info('Saving logs to {:s}'.format(log_file))

    if args.resume is not None:

        ckpt_idx = args.fileid
        ckpt_filename = os.path.join(args.resume, 'checkpoint_{:04d}.pth.tar'.format(ckpt_idx))
        assert os.path.isfile(ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename, map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        logger.info("=> loading checkpoint '{}', current iou: {:.04f}".format(ckpt_filename, train_iou))


    model = cuda_model.convertModel2Cuda(model, gpu_id=args.gpu_id, multiGpu=args.multiGpu)
    train_dataset = cDataset(dataset_split='train', seq_length=args.seq_len, sample_rate=[4], rdOffset=True, rdDrop=True)
    val_dataset =cDataset(dataset_split='val', seq_length=args.seq_len, sample_rate=[4], rdDrop=False, rdOffset=False)


    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=False,
                                  num_workers=4)

    model_optim = optim.Adam(filter(lambda p:p.requires_grad,  model.parameters()), lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim, 'min', patience=10)

    cls_weights = torch.FloatTensor([0.05, 1.0]).cuda()
    # cls_weights = None

    widgets = ['Test: ', ' -- [ ', progressbar.Counter(), '|', str(len(train_dataloader)), ' ] ',
               progressbar.Bar(), ' cls loss:  ', progressbar.FormatLabel(''),
               ' loc loss: ', progressbar.FormatLabel(''),
               ' IoU : ', progressbar.FormatLabel(''),
               ' (', progressbar.ETA(), ' ) ']

    # bar = progressbar.ProgressBar(max_value=step_per_epoch, widgets=widgets)
    # bar.start()

    for epoch in range(args.start_epoch, args.nof_epoch+args.start_epoch):

            total_losses = AverageMeter()
            loc_losses = AverageMeter()
            cls_losses = AverageMeter()
            matched_IOU = AverageMeter()
            true_IOU = AverageMeter()
            model.train()
            pbar = progressbar.ProgressBar(max_value=len(train_dataloader), widgets=widgets)
            pbar.start()
            for i_batch, sample_batch in enumerate(train_dataloader):
                # pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                start_indices = Variable(sample_batch[1])
                end_indices = Variable(sample_batch[2])
                gt_valids = Variable(sample_batch[3])
                # gt_overlaps = Variable(sample_batch[4])

                # seq_labels = Variable(sample_batch[3])

                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    start_indices = start_indices.cuda()
                    end_indices = end_indices.cuda()

                gt_positions = torch.stack([start_indices, end_indices], dim=-1)


                head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(feature_batch)

                pred_positions = torch.stack([head_positions, tail_positions], dim=-1)
                # pred_scores = F.sigmoid(cls_scores)

                if args.hmatch:
                    assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_v2(gt_positions, pred_positions, gt_valids, thres=args.hassign_thres)
                else:
                    #FIXME: do it later!
                    assigned_scores, assigned_locations, total_valid, total_iou = f_match.Assign_Batch_v2(gt_positions, pred_positions, gt_valids, thres=args.hassign_thres)
                    # _, _, total_valid, total_iou = h_match.Assign_Batch_v2(gt_positions, pred_positions, gt_valids, thres=args.hassign_thres)
                true_valid, true_iou = h_match.totalMatch_Batch(gt_positions, pred_positions, gt_valids)
                assert true_valid == total_valid, 'WRONG'

                if total_valid>0:
                    matched_IOU.update(total_iou / total_valid, total_valid)
                    true_IOU.update(true_iou/total_valid, total_valid)

                assigned_scores = Variable(torch.LongTensor(assigned_scores),requires_grad=False)
                # assigned_overlaps = Variable(torch.FloatTensor(assigned_overlaps), requires_grad=False)
                assigned_locations = Variable(torch.LongTensor(assigned_locations), requires_grad=False)
                if use_cuda:
                    assigned_scores = assigned_scores.cuda()
                    assigned_locations = assigned_locations.cuda()
                    # assigned_overlaps = assigned_overlaps.cuda()

                # pred_scores = pred_scores.contiguous().view(-1)
                # assigned_scores = assigned_scores.contiguous().view(-1)
                # assigned_overlaps = assigned_overlaps.contiguous().view(-1)
                # cls_loss = ClsLocLoss2_OneClsRegression(pred_scores, assigned_scores, assigned_overlaps)
                cls_scores = cls_scores.contiguous().view(-1, cls_scores.size()[-1])
                assigned_scores = assigned_scores.contiguous().view(-1)

                cls_loss = F.cross_entropy(cls_scores, assigned_scores, weight=cls_weights)

                if total_valid>0:
                    assigned_head_positions = assigned_locations[:,:,0]
                    assigned_head_positions = assigned_head_positions.contiguous().view(-1)
                    #
                    assigned_tail_positions = assigned_locations[:,:,1]
                    assigned_tail_positions = assigned_tail_positions.contiguous().view(-1)


                    head_pointer_probs = head_pointer_probs.contiguous().view(-1, head_pointer_probs.size()[-1])
                    tail_pointer_probs = tail_pointer_probs.contiguous().view(-1, tail_pointer_probs.size()[-1])

                    assigned_head_positions = torch.masked_select(assigned_head_positions, assigned_scores.byte())
                    assigned_tail_positions = torch.masked_select(assigned_tail_positions, assigned_scores.byte())

                    head_pointer_probs = torch.index_select(head_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))
                    tail_pointer_probs = torch.index_select(tail_pointer_probs, dim=0, index=assigned_scores.nonzero().squeeze(1))

                    # if args.EMD:
                    assigned_head_positions = to_one_hot(assigned_head_positions, args.seq_len)
                    assigned_tail_positions = to_one_hot(assigned_tail_positions, args.seq_len)

                    prediction_head_loss = Simple_L2(head_pointer_probs, assigned_head_positions, needSoftMax=True)
                    prediction_tail_loss = Simple_L2(tail_pointer_probs, assigned_tail_positions, needSoftMax=True)
                    # else:
                    #     prediction_head_loss = F.cross_entropy(head_pointer_probs, assigned_head_positions)
                    #     prediction_tail_loss = F.cross_entropy(tail_pointer_probs, assigned_tail_positions)
                    loc_losses.update(prediction_head_loss.data.item() + prediction_tail_loss.data.item(),
                                      total_valid)#FIXME
                    total_loss =  args.alpha*(prediction_head_loss + prediction_tail_loss) + cls_loss
                else:
                    total_loss = cls_loss

                model_optim.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
                model_optim.step()
                cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
                total_losses.update(total_loss.item(), feature_batch.size(0))

                widgets[-8] = progressbar.FormatLabel('{:04.4f}'.format(cls_losses.avg))
                widgets[-6] = progressbar.FormatLabel('{:04.4f}'.format(loc_losses.avg))
                widgets[-4] = progressbar.FormatLabel('{:01.4f}'.format(matched_IOU.avg))
                pbar.update(i_batch)


            logger.info(
                "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tAvg-matched_IOU:{:.4f}\t Avg-true-IOU:{:.4f}".format(
                    epoch,
                    model_optim.param_groups[0]['lr'], total_losses.avg, cls_losses.avg, loc_losses.avg, matched_IOU.avg, true_IOU.avg))
            train_iou = matched_IOU.avg
            optim_scheduler.step(total_losses.avg)

            model.eval()

            matched_IOU = AverageMeter()
            pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
            for i_batch, sample_batch in enumerate(val_dataloader):
                pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                start_indices = Variable(sample_batch[1])
                end_indices = Variable(sample_batch[2])
                gt_valids = Variable(sample_batch[3])
                # valid_indices = Variable(sample_batch[3])

                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    start_indices = start_indices.cuda()
                    end_indices = end_indices.cuda()

                gt_positions = torch.stack([start_indices, end_indices], dim=-1)

                head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                    feature_batch)

                pred_positions = torch.stack([head_positions, tail_positions], dim=-1)

                # assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_eval(gt_positions, pred_positions, gt_valids, thres=args.hassign_thres) #FIXME
                matched_valid, matched_iou = h_match.totalMatch_Batch(gt_positions, pred_positions, gt_valids)
                if matched_valid>0:
                    matched_IOU.update(matched_iou / matched_valid, matched_valid)

            logger.info(
                "Val -- Epoch :{:06d}, LR: {:.6f},\tloc-Avg-matched_IOU:{:.4f}".format(
                    epoch,model_optim.param_groups[0]['lr'], matched_IOU.avg, ))


            if epoch % 1 == 0 :
                save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'loss':total_losses.avg,
            'cls_loss': cls_losses.avg,
            'loc_loss': loc_losses.avg,
            'train-IOU':train_iou,
            'IoU': matched_IOU.avg}, (epoch+1), file_direcotry=save_directory)
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = dir_utils.get_dir(
        os.path.join(
            project_root, 'ckpts',
            'Delete-{:s}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-seqlen{:d}-{:s}-{:s}'
            .format(script_name_stem, args.hassign_thres, args.alpha,
                    args.hidden_dim, args.dropout, args.seq_len,
                    loss_type[args.EMD], match_type[args.hmatch])))
    log_file = os.path.join(save_directory,
                            'log-{:s}.txt'.format(dir_utils.get_date_str()))
    logger = chinese_utils.get_logger(log_file)
    chinese_utils.print_config(vars(args), logger)

    model = PointerNetwork(input_dim=args.input_dim,
                           embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim,
                           max_decoding_len=args.net_outputs,
                           dropout=args.dropout,
                           n_enc_layers=2)
    hassign_thres = args.hassign_thres
    logger.info("Number of Params\t{:d}".format(
        sum([p.data.nelement() for p in model.parameters()])))
    logger.info('Saving logs to {:s}'.format(log_file))

    if args.resume is not None:

        ckpt_idx = 48

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(
            ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename,
                                map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        logger.info("=> loading checkpoint '{}', current iou: {:.04f}".format(
            ckpt_filename, train_iou))

    model = cuda_model.convertModel2Cuda(model,
                                         gpu_id=args.gpu_id,
                                         multiGpu=args.multiGpu)

    train_dataset = cDataset(dataset_split='train')
    val_dataset = cDataset(dataset_split='val')

    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False,
                                num_workers=4)

    model_optim = optim.Adam(filter(lambda p: p.requires_grad,
                                    model.parameters()),
                             lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim,
                                                           'min',
                                                           patience=10)

    alpha = args.alpha
    # cls_weights = torch.FloatTensor([0.05, 1.0]).cuda()
    for epoch in range(args.start_epoch, args.nof_epoch + args.start_epoch):
        total_losses = AverageMeter()
        loc_losses = AverageMeter()
        cls_losses = AverageMeter()
        Accuracy = AverageMeter()
        IOU = AverageMeter()
        ordered_IOU = AverageMeter()
        model.train()
        pbar = progressbar.ProgressBar(max_value=len(train_dataloader))
        for i_batch, sample_batch in enumerate(train_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            start_indices = Variable(sample_batch[1])
            end_indices = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])
            # seq_labels = Variable(sample_batch[4])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                start_indices = start_indices.cuda()
                end_indices = end_indices.cuda()

            gt_positions = torch.stack([start_indices, end_indices], dim=-1)

            head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                feature_batch)

            pred_positions = torch.stack([head_positions, tail_positions],
                                         dim=-1)
            if args.hmatch:
                assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_v2(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)
                IOU.update(total_iou / total_valid, total_valid)
            else:
                assigned_scores, assigned_locations = f_match.Assign_Batch(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)
                _, _, total_valid, total_iou = h_match.Assign_Batch_v2(
                    gt_positions,
                    pred_positions,
                    gt_valids,
                    thres=hassign_thres)
                IOU.update(total_iou / total_valid, total_valid)

            assigned_scores = Variable(torch.LongTensor(assigned_scores),
                                       requires_grad=False)
            assigned_locations = Variable(torch.LongTensor(assigned_locations),
                                          requires_grad=False)
            if use_cuda:
                assigned_scores = assigned_scores.cuda()
                assigned_locations = assigned_locations.cuda()

            cls_scores = cls_scores.contiguous().view(-1,
                                                      cls_scores.size()[-1])
            assigned_scores = assigned_scores.contiguous().view(-1)

            cls_loss = F.cross_entropy(cls_scores, assigned_scores)

            if total_valid > 0:
                assigned_head_positions = assigned_locations[:, :, 0]
                assigned_head_positions = assigned_head_positions.contiguous(
                ).view(-1)
                #
                assigned_tail_positions = assigned_locations[:, :, 1]
                assigned_tail_positions = assigned_tail_positions.contiguous(
                ).view(-1)

                head_pointer_probs = head_pointer_probs.contiguous().view(
                    -1,
                    head_pointer_probs.size()[-1])
                tail_pointer_probs = tail_pointer_probs.contiguous().view(
                    -1,
                    tail_pointer_probs.size()[-1])

                assigned_head_positions = torch.masked_select(
                    assigned_head_positions, assigned_scores.byte())
                assigned_tail_positions = torch.masked_select(
                    assigned_tail_positions, assigned_scores.byte())

                head_pointer_probs = torch.index_select(
                    head_pointer_probs,
                    dim=0,
                    index=assigned_scores.nonzero().squeeze(1))
                tail_pointer_probs = torch.index_select(
                    tail_pointer_probs,
                    dim=0,
                    index=assigned_scores.nonzero().squeeze(1))

                if args.EMD:
                    assigned_head_positions = to_one_hot(
                        assigned_head_positions, args.seq_len)
                    assigned_tail_positions = to_one_hot(
                        assigned_tail_positions, args.seq_len)

                    prediction_head_loss = EMD_L2(head_pointer_probs,
                                                  assigned_head_positions,
                                                  needSoftMax=True)
                    prediction_tail_loss = EMD_L2(tail_pointer_probs,
                                                  assigned_tail_positions,
                                                  needSoftMax=True)
                else:
                    prediction_head_loss = F.cross_entropy(
                        head_pointer_probs, assigned_head_positions)
                    prediction_tail_loss = F.cross_entropy(
                        tail_pointer_probs, assigned_tail_positions)
                loc_losses.update(
                    prediction_head_loss.data.item() +
                    prediction_tail_loss.data.item(), feature_batch.size(0))
                total_loss = alpha * (prediction_head_loss +
                                      prediction_tail_loss) + cls_loss
            else:
                total_loss = cls_loss

            model_optim.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
            model_optim.step()
            cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
            total_losses.update(total_loss.item(), feature_batch.size(0))

        logger.info(
            "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tcls-Accuracy:{:.4f}\tloc-Avg-IOU:{:.4f}\t topIOU:{:.4f}"
            .format(epoch, model_optim.param_groups[0]['lr'], total_losses.avg,
                    cls_losses.avg, loc_losses.avg, Accuracy.avg, IOU.avg,
                    ordered_IOU.avg))

        optim_scheduler.step(total_losses.avg)

        model.eval()

        IOU = AverageMeter()
        pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
        for i_batch, sample_batch in enumerate(val_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            start_indices = Variable(sample_batch[1])
            end_indices = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])
            # valid_indices = Variable(sample_batch[4])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                start_indices = start_indices.cuda()
                end_indices = end_indices.cuda()

            gt_positions = torch.stack([start_indices, end_indices], dim=-1)

            head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                feature_batch
            )  #Update: compared to the previous version, we now update the matching rules

            pred_positions = torch.stack([head_positions, tail_positions],
                                         dim=-1)

            #TODO: should NOT change here for evaluation!
            assigned_scores, assigned_locations, total_valid, total_iou = h_match.Assign_Batch_eval(
                gt_positions, pred_positions, gt_valids, thres=hassign_thres)
            IOU.update(total_iou / total_valid, total_valid)

        logger.info(
            "Val -- Epoch :{:06d}, LR: {:.6f},\tloc-Avg-IOU:{:.4f}".format(
                epoch,
                model_optim.param_groups[0]['lr'],
                IOU.avg,
            ))

        if epoch % 1 == 0:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'loss': total_losses.avg,
                    'cls_loss': cls_losses.avg,
                    'loc_loss': loc_losses.avg,
                    'IoU': IOU.avg
                }, (epoch + 1),
                file_direcotry=save_directory)
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)

    # Pretty print the run args
    pp.pprint(vars(args))

    model = PointerNetwork(input_dim=args.input_dim,
                           embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim,
                           max_decoding_len=args.net_outputs,
                           dropout=args.dropout,
                           n_enc_layers=2)
    hassign_thres = args.hassign_thres
    print("Number of Params\t{:d}".format(
        sum([p.data.nelement() for p in model.parameters()])))
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = '{:s}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-ckpt'.format(
        script_name_stem, hassign_thres, args.alpha, args.hidden_dim,
        args.dropout)
    print("Save ckpt to {:s}".format(save_directory))
    if args.resume is not None:

        ckpt_idx = 7

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(
            ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename,
                                map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        print("=> loading checkpoint '{}', current iou: {:.04f}".format(
            ckpt_filename, train_iou))

    model = cuda_model.convertModel2Cuda(model,
                                         gpu_id=args.gpu_id,
                                         multiGpu=args.multiGpu)

    # train_dataset  = THUMOST14(seq_length=args.seq_len, overlap=0.9, sample_rate=[4], dataset_split='train',rdDrop=True,rdOffset=True)

    val_dataset = THUMOST14(seq_length=args.seq_len,
                            overlap=0.9,
                            sample_rate=[4],
                            dataset_split='val',
                            rdDrop=False,
                            rdOffset=False)

    # train_dataloader = DataLoader(train_dataset,
    #                               batch_size=args.batch_size,
    #                               shuffle=True,
    #                               num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False,
                                num_workers=4)

    model_optim = optim.Adam(filter(lambda p: p.requires_grad,
                                    model.parameters()),
                             lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim,
                                                           'min',
                                                           patience=20)

    alpha = args.alpha
    cls_weights = torch.FloatTensor([0.05, 1.0]).cuda()
    for epoch in range(args.start_epoch, args.nof_epoch + args.start_epoch):
        total_losses = AverageMeter()
        loc_losses = AverageMeter()
        cls_losses = AverageMeter()
        Accuracy = AverageMeter()
        IOU = AverageMeter()
        ordered_IOU = AverageMeter()
        model.train()
        pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
        for i_batch, sample_batch in enumerate(val_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            start_indices = Variable(sample_batch[1])
            end_indices = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                start_indices = start_indices.cuda()
                end_indices = end_indices.cuda()

            gt_positions = torch.stack([start_indices, end_indices], dim=-1)

            head_pointer_probs, head_positions, tail_pointer_probs, tail_positions, cls_scores, _ = model(
                feature_batch)

            pred_positions = torch.stack([head_positions, tail_positions],
                                         dim=-1)

            assigned_scores, assigned_locations = h_assign.Assign_Batch(
                gt_positions, pred_positions, gt_valids, thres=hassign_thres)
            if np.sum(assigned_scores) > 0:
                print "Output at {:d}".format(i_batch)
                # n_valid = valid_indices.data[0, 0]
                # view_idx = valid_indices.nonzero()[0][0].item()
                # n_valid = valid_indices[view_idx, 0].item()
                print "GT:"
                print(assigned_locations[0])
                print("Pred")
                print(pred_positions[0])
                _, head_sort = head_pointer_probs[0, 0, :].sort()
                _, tail_sort = tail_pointer_probs[0, 0, :].sort()

                print("END of {:d}".format(i_batch))
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = dir_utils.get_dir(os.path.join(project_root, 'ckpts', '{:s}-assgin{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}-seqlen{:d}-{:s}-{:s}'.
                                  format(script_name_stem, args.hassign_thres, args.alpha, args.hidden_dim, args.dropout, args.seq_len, loss_type[args.EMD], match_type[args.hmatch])))
    log_file = os.path.join(save_directory, 'log-{:s}.txt'.format(dir_utils.get_date_str()))
    logger = chinese_utils.get_logger(log_file)
    chinese_utils.print_config(vars(args), logger)


    model = BaseLSTMNetwork(input_dim=args.input_dim, embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim, max_decoding_len=args.net_outputs, dropout=args.dropout, n_enc_layers=2)
    hassign_thres = args.hassign_thres
    logger.info("Number of Params\t{:d}".format(sum([p.data.nelement() for p in model.parameters()])))
    logger.info('Saving logs to {:s}'.format(log_file))

    if args.resume is not None:

        ckpt_idx = 48

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename, map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        logger.info("=> loading checkpoint '{}', current iou: {:.04f}".format(ckpt_filename, train_iou))


    model = cuda_model.convertModel2Cuda(model, gpu_id=args.gpu_id, multiGpu=args.multiGpu)

    train_dataset = MNIST(dataset_split='train')
    val_dataset =MNIST(dataset_split='val')


    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=False,
                                  num_workers=4)

    model_optim = optim.Adam(filter(lambda p:p.requires_grad,  model.parameters()), lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim, 'min', patience=10)

    alpha=args.alpha
    # cls_weights = torch.FloatTensor([0.05, 1.0]).cuda()
    for epoch in range(args.start_epoch, args.nof_epoch+args.start_epoch):
            total_losses = AverageMeter()
            loc_losses = AverageMeter()
            cls_losses = AverageMeter()
            Accuracy = AverageMeter()
            IOU = AverageMeter()
            ordered_IOU = AverageMeter()
            model.train()
            pbar = progressbar.ProgressBar(max_value=len(train_dataloader))
            for i_batch, sample_batch in enumerate(train_dataloader):
                pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                labels = Variable(sample_batch[1])


                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    labels = labels.cuda()
                    # end_indices = end_indices.cuda()

                pred_labels = model(feature_batch)

                labels = labels.contiguous().view(-1)
                pred_labels = pred_labels.contiguous().view(-1, pred_labels.size()[-1])

                pred_probs = F.softmax(pred_labels, dim=1)[:, 1]
                pred_probs[pred_probs>0.5] = 1
                pred_probs[pred_probs<=0.5] = -1
                n_positives = torch.sum(labels).item()
                iou = torch.sum(pred_probs==labels.float()).item()*1. / n_positives
                IOU.update(iou, 1.)

                total_loss = F.cross_entropy(pred_labels, labels)

                model_optim.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
                model_optim.step()
                # cls_losses.update(cls_loss.data.item(), feature_batch.size(0))
                total_losses.update(total_loss.item(), feature_batch.size(0))


            logger.info(
                "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t c-loss:{:.4f}, \tloc-loss:{:.4f}\tcls-Accuracy:{:.4f}\tloc-Avg-IOU:{:.4f}\t topIOU:{:.4f}".format(
                    epoch,
                    model_optim.param_groups[0]['lr'], total_losses.avg, cls_losses.avg, loc_losses.avg, Accuracy.avg, IOU.avg, ordered_IOU.avg))

            optim_scheduler.step(total_losses.avg)

            model.eval()

            IOU = AverageMeter()
            pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
            for i_batch, sample_batch in enumerate(val_dataloader):
                pbar.update(i_batch)

                feature_batch = Variable(sample_batch[0])
                labels = Variable(sample_batch[1])

                if use_cuda:
                    feature_batch = feature_batch.cuda()
                    labels = labels.cuda()

                labels = labels.contiguous().view(-1)

                pred_labels = model(feature_batch)

                pred_labels = pred_labels.contiguous().view(-1, pred_labels.size()[-1])

                pred_probs = F.softmax(pred_labels, dim=1)[:, 1]
                n_positives = torch.sum(labels).item()
                pred_probs[pred_probs > 0.5] = 1
                pred_probs[pred_probs <= 0.5] = -1

                iou = torch.sum(pred_probs == labels.float()).item() * 1. / n_positives
                IOU.update(iou, 1.)


            logger.info(
                "Val -- Epoch :{:06d}, LR: {:.6f},\tloc-Avg-IOU:{:.4f}".format(
                    epoch,model_optim.param_groups[0]['lr'], IOU.avg, ))


            if epoch % 1 == 0:
                save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'loss':total_losses.avg,
            'cls_loss': cls_losses.avg,
            'loc_loss': loc_losses.avg,
            'IoU': IOU.avg}, (epoch+1), file_direcotry=save_directory)
def main():
    global args
    args = (parser.parse_args())
    use_cuda = cuda_model.ifUseCuda(args.gpu_id, args.multiGpu)
    script_name_stem = dir_utils.get_stem(__file__)
    save_directory = dir_utils.get_dir(
        os.path.join(
            project_root, 'ckpts',
            '{:s}-{:s}-{:s}-split-{:d}-decoderatio-{:.2f}-alpha{:.4f}-dim{:d}-dropout{:.4f}'
            .format(script_name_stem, args.dataset, args.eval_metrics,
                    args.split, args.decode_ratio, args.alpha, args.hidden_dim,
                    args.dropout)))
    log_file = os.path.join(save_directory,
                            'log-{:s}.txt'.format(dir_utils.get_date_str()))
    logger = log_utils.get_logger(log_file)
    log_utils.print_config(vars(args), logger)

    # get train/val split
    if args.dataset == 'SumMe':
        train_val_perms = np.arange(25)
    elif args.dataset == 'TVSum':
        train_val_perms = np.arange(50)
    # fixed permutation
    random.Random(0).shuffle(train_val_perms)
    train_val_perms = train_val_perms.reshape([5, -1])
    train_perms = np.delete(train_val_perms, args.split, 0).reshape([-1])
    val_perms = train_val_perms[args.split]
    logger.info(" training split: " + str(train_perms))
    logger.info(" val split: " + str(val_perms))

    if args.location == 'home':
        data_path = os.path.join(os.path.expanduser('~'), 'datasets')
    else:
        data_path = os.path.join('/nfs/%s/boyu/SDN' % (args.location),
                                 'datasets')
    train_dataset = vsTVSum_Loader3_c3dd_segment.cDataset(
        dataset_name=args.dataset,
        split='train',
        decode_ratio=args.decode_ratio,
        train_val_perms=train_perms,
        data_path=data_path)
    max_input_len = train_dataset.max_input_len
    maximum_outputs = int(args.decode_ratio * max_input_len)
    val_dataset = vsTVSum_Loader3_c3dd_segment.cDataset(
        dataset_name=args.dataset,
        split='val',
        decode_ratio=args.decode_ratio,
        train_val_perms=val_perms,
        data_path=data_path)
    train_evaluator = Evaluator.Evaluator(dataset_name=args.dataset,
                                          split='tr',
                                          max_input_len=max_input_len,
                                          maximum_outputs=maximum_outputs,
                                          sum_budget=0.15,
                                          train_val_perms=train_perms,
                                          eval_metrics=args.eval_metrics,
                                          data_path=data_path)
    val_evaluator = Evaluator.Evaluator(dataset_name=args.dataset,
                                        split='val',
                                        max_input_len=max_input_len,
                                        maximum_outputs=maximum_outputs,
                                        sum_budget=0.15,
                                        train_val_perms=val_perms,
                                        eval_metrics=args.eval_metrics,
                                        data_path=data_path)

    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=4)
    val_dataloader = DataLoader(val_dataset,
                                batch_size=args.batch_size,
                                shuffle=False,
                                num_workers=4)

    model = PointerNetwork(input_dim=args.input_dim,
                           embedding_dim=args.embedding_dim,
                           hidden_dim=args.hidden_dim,
                           max_decoding_len=maximum_outputs,
                           dropout=args.dropout,
                           n_enc_layers=2,
                           output_classes=1)
    # hassign_thres = args.hassign_thres
    logger.info("Number of Params\t{:d}".format(
        sum([p.data.nelement() for p in model.parameters()])))
    logger.info('Saving logs to {:s}'.format(log_file))

    if args.resume is not None:

        ckpt_idx = 48

        ckpt_filename = args.resume.format(ckpt_idx)
        assert os.path.isfile(
            ckpt_filename), 'Error: no checkpoint directory found!'

        checkpoint = torch.load(ckpt_filename,
                                map_location=lambda storage, loc: storage)
        model.load_state_dict(checkpoint['state_dict'], strict=False)
        train_iou = checkpoint['IoU']
        args.start_epoch = checkpoint['epoch']

        logger.info("=> loading checkpoint '{}', current iou: {:.04f}".format(
            ckpt_filename, train_iou))

    model = cuda_model.convertModel2Cuda(model,
                                         gpu_id=args.gpu_id,
                                         multiGpu=args.multiGpu)

    model_optim = optim.Adam(filter(lambda p: p.requires_grad,
                                    model.parameters()),
                             lr=float(args.lr))
    optim_scheduler = optim.lr_scheduler.ReduceLROnPlateau(model_optim,
                                                           'min',
                                                           patience=10)

    alpha = args.alpha
    # cls_weights = torch.FloatTensor([0.2, 1.0]).cuda()
    # if args.set_cls_weight:
    #     cls_weights = torch.FloatTensor([1.*train_dataset.n_positive_train_samples/train_dataset.n_total_train_samples, args.cls_pos_weight]).cuda()
    # else:
    #     cls_weights = torch.FloatTensor([0.5, 0.5]).cuda()
    # logger.info(" total: {:d}, total pos: {:d}".format(train_dataset.n_total_train_samples, train_dataset.n_positive_train_samples))
    # logger.info(" classify weight: " + str(cls_weights[0]) + str(cls_weights[1]))
    for epoch in range(args.start_epoch, args.nof_epoch + args.start_epoch):
        total_losses = AverageMeter()
        # loc_losses = AverageMeter()
        pointer_losses = AverageMeter()
        rgs_losses = AverageMeter()
        Accuracy = AverageMeter()
        # IOU = AverageMeter()
        # ordered_IOU = AverageMeter()
        model.train()
        pbar = progressbar.ProgressBar(max_value=len(train_dataloader))
        for i_batch, sample_batch in enumerate(train_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            pointer_indices = Variable(sample_batch[1])
            pointer_scores = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])
            # seq_labels = Variable(sample_batch[3])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                pointer_indices = pointer_indices.cuda()
                pointer_scores = pointer_scores.cuda()

            gt_positions = pointer_indices
            gt_scores = pointer_scores

            pointer_probs, pointer_positions, cls_scores, _ = model(
                feature_batch)

            pred_positions = pointer_positions

            cls_scores = cls_scores.contiguous().squeeze(2)

            # print(pointer_probs.size())
            # print(gt_positions.size())
            pointer_loss = F.cross_entropy(pointer_probs.permute(0, 2, 1),
                                           gt_positions)

            rgs_loss = F.mse_loss(cls_scores, gt_scores)

            total_loss = alpha * pointer_loss + rgs_loss

            model_optim.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
            model_optim.step()
            pointer_losses.update(pointer_loss.data.item(),
                                  feature_batch.size(0))
            rgs_losses.update(rgs_loss.data.item(), feature_batch.size(0))
            total_losses.update(total_loss.item(), feature_batch.size(0))

        logger.info(
            "Train -- Epoch :{:06d}, LR: {:.6f},\tloss={:.4f}, \t pointer-loss:{:.4f}, \tregress-loss:{:.4f}\tcls-Accuracy:{:.4f}"
            .format(epoch, model_optim.param_groups[0]['lr'], total_losses.avg,
                    pointer_losses.avg, rgs_losses.avg, Accuracy.avg))

        optim_scheduler.step(total_losses.avg)

        model.eval()

        pointer_losses = AverageMeter()
        rgs_losses = AverageMeter()
        pbar = progressbar.ProgressBar(max_value=len(val_dataloader))
        for i_batch, sample_batch in enumerate(val_dataloader):
            pbar.update(i_batch)

            feature_batch = Variable(sample_batch[0])
            pointer_indices = Variable(sample_batch[1])
            pointer_scores = Variable(sample_batch[2])
            gt_valids = Variable(sample_batch[3])
            # valid_indices = Variable(sample_batch[3])

            if use_cuda:
                feature_batch = feature_batch.cuda()
                pointer_indices = pointer_indices.cuda()
                pointer_scores = pointer_scores.cuda()

            gt_positions = pointer_indices
            gt_scores = pointer_scores

            pointer_probs, pointer_positions, cls_scores, _ = model(
                feature_batch)

            pred_positions = pointer_positions

            cls_scores = cls_scores.contiguous().squeeze(2)

            pointer_loss = F.cross_entropy(pointer_probs.permute(0, 2, 1),
                                           gt_positions)

            rgs_loss = F.mse_loss(cls_scores, gt_scores)

            pointer_losses.update(pointer_loss.data.item(),
                                  feature_batch.size(0))
            rgs_losses.update(rgs_loss.data.item(), feature_batch.size(0))

        train_F1s = train_evaluator.Evaluate(model)
        val_F1s = val_evaluator.Evaluate(model)

        logger.info("Train -- Epoch :{:06d},\tF1s:{:.4f}".format(
            epoch, train_F1s))

        logger.info(
            "Val -- Epoch :{:06d},\t pointer-loss:{:.4f}, \tregress-loss:{:.4f}, \tF1s{:.4f}"
            .format(epoch, pointer_losses.avg, rgs_losses.avg, val_F1s))

        if epoch % 1 == 0:
            save_checkpoint(
                {
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'loss': total_losses.avg,
                    'pointer_loss': pointer_losses.avg,
                    'rgs_loss': rgs_losses.avg,
                    'val_F1s': val_F1s
                }, (epoch + 1),
                file_direcotry=save_directory)