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)
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)
def prepare_env(cfg):
# create output directory
if cfg.BASIC.CREATE_OUTPUT_DIR:
out_dir = cfg.TRAIN.OUTPUT_DIR
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# create directory for prediction
out_dir = cfg.TEST.RESULT_DIR
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# backup codes
if cfg.BASIC.BACKUP_CODES:
backup_dir = cfg.BASIC.BACKUP_DIR
backup_codes(cfg, cfg.BASIC.ROOT_DIR, backup_dir, cfg.BASIC.BACKUP_LIST)