示例#1
0
def main():
    args = parse_args()
    update_config(args.cfg)

    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    cudnn.enabled = cfg.CUDNN.ENABLE

    # data loader
    val_dset = TALDataset(cfg, cfg.DATASET.VAL_SPLIT)
    val_loader = DataLoader(val_dset,
                            batch_size=cfg.TEST.BATCH_SIZE,
                            shuffle=False,
                            drop_last=False,
                            num_workers=cfg.BASIC.WORKERS,
                            pin_memory=cfg.DATASET.PIN_MEMORY)

    model = LocNet(cfg)
    model.cuda()

    #evaluate existing model
    checkpoint = torch.load(args.checkpoint)
    model.load_state_dict(checkpoint['model'])
    epoch = checkpoint['epoch']

    out_df = evaluation(val_loader, model, epoch, cfg)
    final_result_process(out_df, epoch, cfg, flag=0)
示例#2
0
文件: main.py 项目: yilonghe/A2Net
def main():
    args = parse_args()
    update_config(args.cfg)
    fix_random_seed(cfg.SEED)

    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    cudnn.enabled = cfg.CUDNN.ENABLE

    # prepare output directory
    output_dir = os.path.join(cfg.ROOT_DIR, cfg.TRAIN.MODEL_DIR)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # copy config ssad.yaml file to output directory
    cfg_file = os.path.join(output_dir, args.cfg.split('/')[-1])
    shutil.copyfile(args.cfg, cfg_file)

    # data loader
    # Notice: we discard the last data
    train_dset = TALDataset(cfg, cfg.DATASET.TRAIN_SPLIT)
    train_loader = DataLoader(train_dset,
                              batch_size=cfg.TRAIN.BATCH_SIZE,
                              shuffle=True,
                              drop_last=False,
                              num_workers=cfg.WORKERS,
                              pin_memory=cfg.PIN_MEMORY)
    val_dset = TALDataset(cfg, cfg.DATASET.VAL_SPLIT)
    val_loader = DataLoader(val_dset,
                            batch_size=cfg.TEST.BATCH_SIZE,
                            shuffle=False,
                            drop_last=False,
                            num_workers=cfg.WORKERS,
                            pin_memory=cfg.PIN_MEMORY)

    model = LocNet(cfg)
    model.apply(weight_init)
    model.cuda()

    optimizer = optim.Adam(model.parameters(), lr=cfg.TRAIN.LR)
    for epoch in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH + 1):
        if epoch in cfg.TRAIN.LR_DECAY_EPOCHS:
            decay_lr(optimizer, factor=cfg.TRAIN.LR_DECAY_FACTOR)

        loss_train = train(cfg, train_loader, model, optimizer)
        print('epoch %d: loss: %f' % (epoch, loss_train))
        with open(os.path.join(cfg.ROOT_DIR, cfg.TRAIN.LOG_FILE), 'a') as f:
            f.write("epoch %d, loss: %.4f\n" % (epoch, loss_train))

        if epoch % cfg.TEST.EVAL_INTERVAL == 0:
            # model
            weight_file = save_model(cfg,
                                     epoch=epoch,
                                     model=model,
                                     optimizer=optimizer)
            out_df_af, out_df_ab = evaluation(val_loader, model, epoch, cfg)
            out_df_ab['conf'] = out_df_ab['conf'] * cfg.TEST.CONCAT_AB
            out_df = pd.concat([out_df_af, out_df_ab])
            post_process(out_df, epoch, cfg, is_soft_nms=False)
示例#3
0
文件: main.py 项目: yilonghe/A2Net
def main():
    args = parse_args()
    update_config(args.cfg)
    # create output directory
    if cfg.BASIC.CREATE_OUTPUT_DIR:
        out_dir = os.path.join(cfg.BASIC.ROOT_DIR, cfg.TRAIN.MODEL_DIR)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
    # copy config file
    if cfg.BASIC.BACKUP_CODES:
        backup_dir = os.path.join(cfg.BASIC.ROOT_DIR, cfg.TRAIN.MODEL_DIR,
                                  'code')
        backup_codes(cfg.BASIC.ROOT_DIR, backup_dir, cfg.BASIC.BACKUP_LISTS)
    fix_random_seed(cfg.BASIC.SEED)
    if cfg.BASIC.SHOW_CFG:
        pprint.pprint(cfg)

    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    cudnn.enabled = cfg.CUDNN.ENABLE

    # data loader
    train_dset = TALDataset(cfg, cfg.DATASET.TRAIN_SPLIT)
    train_loader = DataLoader(train_dset,
                              batch_size=cfg.TRAIN.BATCH_SIZE,
                              shuffle=True,
                              drop_last=False,
                              num_workers=cfg.BASIC.WORKERS,
                              pin_memory=cfg.DATASET.PIN_MEMORY)
    val_dset = TALDataset(cfg, cfg.DATASET.VAL_SPLIT)
    val_loader = DataLoader(val_dset,
                            batch_size=cfg.TEST.BATCH_SIZE,
                            shuffle=False,
                            drop_last=False,
                            num_workers=cfg.BASIC.WORKERS,
                            pin_memory=cfg.DATASET.PIN_MEMORY)

    model = LocNet(cfg)
    model.apply(weight_init)
    model.cuda()

    optimizer = optim.Adam(model.parameters(), lr=cfg.TRAIN.LR)
    for epoch in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH + 1):
        loss_train = train(cfg, train_loader, model, optimizer)
        print('epoch %d: loss: %f' % (epoch, loss_train))
        with open(os.path.join(cfg.BASIC.ROOT_DIR, cfg.TRAIN.LOG_FILE),
                  'a') as f:
            f.write("epoch %d, loss: %.4f\n" % (epoch, loss_train))

        # decay lr
        if epoch in cfg.TRAIN.LR_DECAY_EPOCHS:
            decay_lr(optimizer, factor=cfg.TRAIN.LR_DECAY_FACTOR)

        if epoch in cfg.TEST.EVAL_INTERVAL:
            save_model(cfg, epoch=epoch, model=model, optimizer=optimizer)
            out_df_ab, out_df_af = evaluation(val_loader, model, epoch, cfg)
            out_df_list = [out_df_ab, out_df_af]
            final_result_process(out_df_list, epoch, cfg, flag=0)
示例#4
0
def main():
    args = parse_args()
    update_config(args.cfg)
    # create output directory
    if cfg.BASIC.CREATE_OUTPUT_DIR:
        out_dir = os.path.join(cfg.BASIC.ROOT_DIR, cfg.TRAIN.MODEL_DIR)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
    # copy config file
    if cfg.BASIC.BACKUP_CODES:
        backup_dir = os.path.join(cfg.BASIC.ROOT_DIR, cfg.TRAIN.MODEL_DIR, 'code')
        backup_codes(cfg.BASIC.ROOT_DIR, backup_dir, cfg.BASIC.BACKUP_LISTS)
    fix_random_seed(cfg.BASIC.SEED)
    if cfg.BASIC.SHOW_CFG:
        pprint.pprint(cfg)

    # cudnn related setting
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    cudnn.enabled = cfg.CUDNN.ENABLE

    # data loader
    train_dset = TALDataset(cfg, cfg.DATASET.TRAIN_SPLIT)
    train_loader = DataLoader(train_dset, batch_size=cfg.TRAIN.BATCH_SIZE,
                              shuffle=True, drop_last=False, num_workers=cfg.BASIC.WORKERS, pin_memory=cfg.DATASET.PIN_MEMORY)
    val_dset = TALDataset(cfg, cfg.DATASET.VAL_SPLIT)
    val_loader = DataLoader(val_dset, batch_size=cfg.TEST.BATCH_SIZE,
                            shuffle=False, drop_last=False, num_workers=cfg.BASIC.WORKERS, pin_memory=cfg.DATASET.PIN_MEMORY)

    model = LocNet(cfg)
    model.apply(weight_init)
    model.cuda()

    #evaluate existing model
    weight_file = '/disk3/zt/code/2_TIP_rebuttal/2_A2Net/output/thumos/output_toy/model_100.pth'
    epoch = 4
    checkpoint = torch.load(weight_file)
    model.load_state_dict(checkpoint['model'])
    out_df_ab, out_df_af = evaluation(val_loader, model, epoch, cfg)

    '''
    flag:
    0: jointly consider out_df_ab and out_df_af
    1: only consider out_df_ab
    2: only consider out_df_af    
    '''
    # evaluate both branch
    out_df_list = [out_df_ab, out_df_af]
    final_result_process(out_df_list, epoch, cfg, flag=0)