def BSN_Train_PEM(opt):
writer = SummaryWriter()
model = PEM(opt)
model = torch.nn.DataParallel(model, device_ids=GPU_IDs).cuda()
optimizer = optim.Adam(model.parameters(),lr=opt["pem_training_lr"],weight_decay = opt["pem_weight_decay"])
def collate_fn(batch):
batch_data = torch.cat([x[0] for x in batch])
batch_iou = torch.cat([x[1][0] for x in batch])
batch_is_whole_lenght = torch.cat([x[1][1] for x in batch])
return batch_data,batch_iou, batch_is_whole_lenght
train_loader = torch.utils.data.DataLoader(ProposalDataSet(opt,subset="train"),
batch_size=model.module.batch_size, shuffle=True,
num_workers=8, pin_memory=True,drop_last=True,collate_fn=collate_fn)
test_loader = torch.utils.data.DataLoader(ProposalDataSet(opt,subset="validation"),
batch_size=model.module.batch_size, shuffle=True,
num_workers=8, pin_memory=True,drop_last=True,collate_fn=collate_fn)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size = opt["pem_step_size"], gamma = opt["pem_step_gamma"])
for epoch in range(opt["pem_epoch"]):
scheduler.step()
train_PEM(train_loader,model,optimizer,epoch,writer,opt)
test_PEM(test_loader,model,epoch,writer,opt)
writer.close()
def BSN_inference_PEM(opt):
model = PEM(opt)
checkpoint = torch.load(opt["checkpoint_path"]+"/"+opt["arch"]+"_pem_best.pth.tar")
base_dict = {'.'.join(k.split('.')[1:]): v for k,v in list(checkpoint['state_dict'].items())}
model.load_state_dict(base_dict)
model = torch.nn.DataParallel(model, device_ids=GPU_IDs).cuda()
model.eval()
test_loader = torch.utils.data.DataLoader(ProposalDataSet(opt,subset=opt["pem_inference_subset"]),
batch_size=model.module.batch_size, shuffle=False,
num_workers=8, pin_memory=True,drop_last=False)
for idx,(video_feature,video_xmin,video_xmax,video_xmin_score,video_xmax_score) in enumerate(test_loader):
video_name = test_loader.dataset.video_list[idx]
video_conf = model(video_feature).view(-1).detach().cpu().numpy()
video_xmin = video_xmin.view(-1).cpu().numpy()
video_xmax = video_xmax.view(-1).cpu().numpy()
video_xmin_score = video_xmin_score.view(-1).cpu().numpy()
video_xmax_score = video_xmax_score.view(-1).cpu().numpy()
df=pd.DataFrame()
df["xmin"]=video_xmin
df["xmax"]=video_xmax
df["xmin_score"]=video_xmin_score
df["xmax_score"]=video_xmax_score
df["iou_score"]=video_conf
df.to_csv("./output/"+opt["arch"]+opt["fix_scale"]+"_PEM_results/"+video_name+".csv",index=False)
def BSN_inference_PEM(opt):
'''
step - 1. load the PEM-model
step - 2. load the dataset
'''
# step - 1
model = PEM(opt)
checkpoint = torch.load(opt["checkpoint_path"]+"/pem_best.pth.tar")
base_dict = {'.'.join(k.split('.')[1:]): v for k,v in list(checkpoint['state_dict'].items())}
model.load_state_dict(base_dict)
model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
model.eval()
# pem_inference_subset = 'validation'
test_loader = torch.utils.data.DataLoader(ProposalDataSet(opt, subset=opt["pem_inference_subset"]),
batch_size=model.module.batch_size, shuffle=False,
num_workers=8, pin_memory=True, drop_last=False)
count = 0
for idx, (video_feature, video_xmin, video_xmax, video_xmin_score, video_xmax_score) in enumerate(test_loader):
video_name = test_loader.dataset.video_list[idx]
video_conf = model(video_feature).view(-1).detach().cpu().numpy() # prob for proposal
video_xmin = video_xmin.view(-1).cpu().numpy()
video_xmax = video_xmax.view(-1).cpu().numpy()
video_xmin_score = video_xmin_score.view(-1).cpu().numpy()
video_xmax_score = video_xmax_score.view(-1).cpu().numpy()
df = pd.DataFrame()
df['xmin'] = video_xmin
df['xmax'] = video_xmax
df['xmin_score'] = video_xmin_score
df['xmax_score'] = video_xmax_score
df['iou_score'] = video_conf
df.to_csv("./output/PEM_results/"+video_name+".csv",index=False)
count += 1
print('there are %5d results' % count)
def BSN_inference_PEM(opt):
model = PEM(opt)
def BSN_inference_PEM(opt):
output_dir = os.path.join(opt['pem_inference_results_dir'],
opt['checkpoint_path'].split('/')[-1])
checkpoint_epoch = opt['checkpoint_epoch']
if checkpoint_epoch is not None:
checkpoint_path = os.path.join(
opt['checkpoint_path'], 'pem_checkpoint.%d.pth' % checkpoint_epoch)
output_dir = os.path.join(output_dir, 'ckpt.%d' % checkpoint_epoch)
else:
checkpoint_path = os.path.join(opt['checkpoint_path'], 'pem_best.pth')
output_dir = os.path.join(output_dir, 'ckpt.best')
print('Checkpoint path is ', checkpoint_path)
checkpoint = torch.load(checkpoint_path)
base_dict = {
'.'.join(k.split('.')[1:]): v
for k, v in list(checkpoint['state_dict'].items())
}
model = PEM(opt)
model.load_state_dict(base_dict)
model = torch.nn.DataParallel(model).cuda()
model.eval()
if not os.path.exists(output_dir):
os.makedirs(output_dir)
test_loader = torch.utils.data.DataLoader(
ProposalDataSet(opt, subset=opt["pem_inference_subset"]),
batch_size=model.module.batch_size,
shuffle=False,
num_workers=opt['data_workers'],
pin_memory=True,
drop_last=False)
current_video = None
columns = ["xmin", "xmax", "xmin_score", "xmax_score", "iou_score"]
for idx, (index_list, video_feature, video_xmin, video_xmax,
video_xmin_score, video_xmax_score) in enumerate(test_loader):
video_conf = model(video_feature).view(-1).detach().cpu().numpy()
video_xmin = video_xmin.view(-1).cpu().numpy()
video_xmax = video_xmax.view(-1).cpu().numpy()
video_xmin_score = video_xmin_score.view(-1).cpu().numpy()
video_xmax_score = video_xmax_score.view(-1).cpu().numpy()
index_list = index_list.numpy()
for batch_idx, full_idx in enumerate(index_list):
video, frame = test_loader.dataset.indices[full_idx]
if not current_video:
print('First video: ', video, full_idx)
current_video = video
current_data = [[] for _ in range(len(columns))]
elif video != current_video:
print('Changing from video %s to video %s: %d' %
(current_video, video, full_idx))
video_result = np.stack(current_data, axis=1)
video_df = pd.DataFrame(video_result, columns=columns)
path = os.path.join(output_dir, '%s.csv' % current_video)
video_df.to_csv(path, index=False)
current_video = video
current_data = [[] for _ in range(len(columns))]
current_data[0].append(video_xmin[batch_idx])
current_data[1].append(video_xmax[batch_idx])
current_data[2].append(video_xmin_score[batch_idx])
current_data[3].append(video_xmax_score[batch_idx])
current_data[4].append(video_conf[batch_idx])
if current_data[0]:
video_result = np.stack(current_data, axis=1)
video_df = pd.DataFrame(video_result, columns=columns)
path = os.path.join(output_dir, '%s.csv' % current_video)
video_df.to_csv(path, index=False)
def BSN_Train_PEM(opt):
model = PEM(opt)
model = torch.nn.DataParallel(model).cuda()
optimizer = optim.Adam(model.parameters(),
lr=opt["pem_training_lr"],
weight_decay=opt["pem_weight_decay"])
print('Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
def collate_fn(batch):
batch_data = torch.cat([x[0] for x in batch])
batch_iou = torch.cat([x[1] for x in batch])
return batch_data, batch_iou
train_dataset = ProposalDataSet(opt, subset="train")
train_sampler = ProposalSampler(train_dataset.proposals,
train_dataset.indices,
max_zero_weight=opt['pem_max_zero_weight'])
global_step = 0
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=model.module.batch_size,
shuffle=False,
sampler=train_sampler,
num_workers=opt['data_workers'],
pin_memory=True,
drop_last=False,
collate_fn=collate_fn if not opt['pem_do_index'] else None)
subset = "validation" if opt['dataset'] == 'activitynet' else "test"
test_loader = torch.utils.data.DataLoader(
ProposalDataSet(opt, subset=subset),
batch_size=model.module.batch_size,
shuffle=True,
num_workers=opt['data_workers'],
pin_memory=True,
drop_last=False,
collate_fn=collate_fn if not opt['pem_do_index'] else None)
milestones = [int(k) for k in opt['pem_lr_milestones'].split(',')]
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=milestones, gamma=opt['pem_step_gamma'])
if opt['log_to_comet']:
comet_exp = CometExperiment(api_key="hIXq6lDzWzz24zgKv7RYz6blo",
project_name="bsnpem",
workspace="cinjon",
auto_metric_logging=True,
auto_output_logging=None,
auto_param_logging=False)
elif opt['local_comet_dir']:
comet_exp = OfflineExperiment(api_key="hIXq6lDzWzz24zgKv7RYz6blo",
project_name="bsnpem",
workspace="cinjon",
auto_metric_logging=True,
auto_output_logging=None,
auto_param_logging=False,
offline_directory=opt['local_comet_dir'])
else:
comet_exp = None
if comet_exp:
comet_exp.log_parameters(opt)
comet_exp.set_name(opt['name'])
test_PEM(test_loader, model, -1, -1, comet_exp, opt)
for epoch in range(opt["pem_epoch"]):
global_step = train_PEM(train_loader, model, optimizer, epoch,
global_step, comet_exp, opt)
test_PEM(test_loader, model, epoch, global_step, comet_exp, opt)
scheduler.step()
