def runner_func(dataset, state_dict, stats, gpu_id, index_queue, result_queue,
iou_dict, dis_dict):
torch.cuda.set_device(gpu_id)
net = PGCN(model_configs, graph_configs, dataset_configs, test_mode=True)
net.load_state_dict(state_dict)
# net.prepare_test_fc()
net.eval()
net.cuda()
while True:
index = index_queue.get()
rel_props, prop_ticks, video_id, n_frames = dataset[index]
# calculate scores
# n_out = prop_ticks.size(0)
# act_scores = torch.zeros((n_out, num_class + 1)).cuda()
# comp_scores = torch.zeros((n_out, num_class)).cuda()
# if not args.no_regression:
# reg_scores = torch.zeros((n_out, num_class * 2)).cuda()
# else:
# reg_scores = None
# load prop fts
vid_full_name = video_id
vid = vid_full_name.split('/')[-1]
act_all_fts, comp_all_fts = I3D_Pooling(
prop_ticks, vid, dataset_configs['test_ft_path'], n_frames)
with torch.no_grad():
vid_iou_dict = torch.from_numpy(iou_dict[vid]).cuda()
vid_dis_dict = torch.from_numpy(dis_dict[vid]).cuda()
act_scores, comp_scores, reg_scores = net(act_all_fts.cuda(),
comp_all_fts.cuda(),
iou_dict=vid_iou_dict,
dis_dict=vid_dis_dict)
if reg_scores is not None:
reg_scores = reg_scores.view(-1, num_class, 2)
reg_scores[:, :,
0] = reg_scores[:, :, 0] * stats[1, 0] + stats[0, 0]
reg_scores[:, :,
1] = reg_scores[:, :, 1] * stats[1, 1] + stats[0, 1]
result_queue.put((dataset.video_list[index].id, rel_props.numpy(),
act_scores.squeeze().cpu().numpy(),
comp_scores.squeeze().cpu().numpy(),
reg_scores.cpu().numpy(), 0))
def main():
global args, best_loss, writer, adj_num, logger
configs = get_and_save_args(parser)
parser.set_defaults(**configs)
dataset_configs = configs["dataset_configs"]
model_configs = configs["model_configs"]
graph_configs = configs["graph_configs"]
args = parser.parse_args()
"""copy codes and creat dir for saving models and logs"""
if not os.path.isdir(args.snapshot_pref):
os.makedirs(args.snapshot_pref)
logger = get_logger(args)
logger.info('\ncreating folder: ' + args.snapshot_pref)
# if not args.evaluate:
writer = SummaryWriter(args.snapshot_pref)
# recorder = Recorder(args.snapshot_pref, ["models", "__pycache__"])
# recorder.writeopt(args)
logger.info('\nruntime args\n\n{}\n\nconfig\n\n{}'.format(
args, dataset_configs))
"""construct model"""
model = PGCN(model_configs, graph_configs)
policies = model.get_optim_policies()
model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()
if args.resume:
if os.path.isfile(args.resume):
logger.info(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
model.load_state_dict(checkpoint['state_dict'])
logger.info(("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
logger.info(("=> no checkpoint found at '{}'".format(args.resume)))
"""construct dataset"""
train_loader = torch.utils.data.DataLoader(
PGCNDataSet(
dataset_configs,
graph_configs,
prop_file=dataset_configs['train_prop_file'],
prop_dict_path=dataset_configs['train_dict_path'],
ft_path=dataset_configs['train_ft_path'],
epoch_multiplier=dataset_configs['training_epoch_multiplier'],
test_mode=False),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True) # in training we drop the last incomplete minibatch
val_loader = torch.utils.data.DataLoader(PGCNDataSet(
dataset_configs,
graph_configs,
prop_file=dataset_configs['test_prop_file'],
prop_dict_path=dataset_configs['val_dict_path'],
ft_path=dataset_configs['test_ft_path'],
epoch_multiplier=dataset_configs['testing_epoch_multiplier'],
reg_stats=train_loader.dataset.stats,
test_mode=False),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
"""loss and optimizer"""
activity_criterion = torch.nn.CrossEntropyLoss().cuda()
completeness_criterion = CompletenessLoss().cuda()
regression_criterion = ClassWiseRegressionLoss().cuda()
for group in policies:
logger.info(
('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
group['name'], len(group['params']), group['lr_mult'],
group['decay_mult'])))
optimizer = torch.optim.SGD(policies,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.evaluate:
validate(val_loader, model, activity_criterion, completeness_criterion,
regression_criterion, 0)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, args.lr_steps)
train(train_loader, model, activity_criterion, completeness_criterion,
regression_criterion, optimizer, epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
loss = validate(val_loader, model, activity_criterion,
completeness_criterion, regression_criterion,
(epoch + 1) * len(train_loader))
# remember best validation loss and save checkpoint
is_best = loss < best_loss
best_loss = min(loss, best_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_loss': loss,
'reg_stats': torch.from_numpy(train_loader.dataset.stats)
},
is_best,
epoch,
filename='checkpoint.pth.tar')
writer.close()
def runner_func(dataset, state_dict, stats, gpu_id, index_queue, result_queue,
iou_dict, dis_dict):
torch.cuda.set_device(gpu_id)
net = PGCN(model_configs, graph_configs, test_mode=True)
net.load_state_dict(state_dict)
# net.prepare_test_fc()
net.eval()
net.cuda()
while True:
index = index_queue.get()
rel_props, prop_ticks, video_id, n_frames = dataset[index]
# calculate scores
n_out = prop_ticks.size(0)
act_scores = torch.zeros((n_out, num_class + 1)).cuda()
comp_scores = torch.zeros((n_out, num_class)).cuda()
if not args.no_regression:
reg_scores = torch.zeros((n_out, num_class * 2)).cuda()
else:
reg_scores = None
# load prop fts
vid_full_name = video_id
vid = vid_full_name.split('/')[-1]
act_all_fts, comp_all_fts = I3D_Pooling(
prop_ticks, vid, dataset_configs['test_ft_path'], n_frames)
for prop_idx, prop in enumerate(prop_ticks):
if prop_idx >= n_out:
break
with torch.no_grad():
vid_iou_dict = iou_dict[vid]
vid_dis_dict = dis_dict[vid]
# print(len(vid_iou_dict), len(vid_dis_dict), len(prop_ticks))
selected_idx = get_adjacent_batch(prop_idx, vid_iou_dict,
vid_dis_dict)
selected_idx = torch.from_numpy(np.array(selected_idx))
act_ft = act_all_fts[selected_idx, :]
comp_ft = comp_all_fts[selected_idx, :]
act_batch_var = act_ft.unsqueeze(0).cuda()
comp_batch_var = comp_ft.unsqueeze(0).cuda()
act_scores[prop_idx, :], comp_scores[prop_idx, :], \
reg_scores[prop_idx, :] = net((act_batch_var, comp_batch_var), None, None, None)
if reg_scores is not None:
reg_scores = reg_scores.view(-1, num_class, 2)
reg_scores[:, :,
0] = reg_scores[:, :, 0] * stats[1, 0] + stats[0, 0]
reg_scores[:, :,
1] = reg_scores[:, :, 1] * stats[1, 1] + stats[0, 1]
# perform stpp on scores
result_queue.put((dataset.video_list[index].id, rel_props.numpy(),
act_scores.cpu().numpy(), comp_scores.cpu().numpy(),
reg_scores.cpu().numpy(), 0))
train_dataset = PGCNDataSet(dataset_configs, graph_configs,
prop_file=dataset_configs['train_prop_file'],
prop_dict_path=dataset_configs['train_dict_path'],
ft_path=dataset_configs['train_ft_path'],
epoch_multiplier=dataset_configs['training_epoch_multiplier'],
mode='Train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=4, shuffle=True,collate_fn=collate_fn,
num_workers=4, pin_memory=True, drop_last=True)
"""construct model"""
from pgcn_models import PGCN
act_iou_dict, act_dis_dict = train_dataset.act_iou_dict, train_dataset.act_dis_dict
model = PGCN(model_configs, graph_configs, act_iou_dict, act_dis_dict)
policies = model.get_optim_policies()
model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()
activity_criterion = torch.nn.CrossEntropyLoss().cuda()
completeness_criterion = CompletenessLoss().cuda()
regression_criterion = ClassWiseRegressionLoss().cuda()
for i, (batch_act_fts, batch_comp_fts, batch_props_type,
batch_props_label, batch_reg_target, props_num_list, vid_list) in enumerate(tqdm(train_loader)):
batch_size = batch_act_fts.size(0)
activity_out, activity_target, activity_prop_type, \
completeness_out, completeness_target, \
regression_out, regression_labels, regression_target, ohem_num = model(batch_act_fts, batch_comp_fts, batch_props_type,
batch_props_label, batch_reg_target,
props_num_list, vid_list)