def main(args):
model = HyperRegression(args,
input_width=CPI_WIDTH,
input_height=CPI_HEIGHT)
model = model.cuda()
resume_checkpoint = args.resume_checkpoint
print(f"Resume Path: {resume_checkpoint}")
checkpoint = torch.load(resume_checkpoint)
model_serialize = checkpoint['model']
model.load_state_dict(model_serialize)
model.eval()
save_path = os.path.join(os.path.split(resume_checkpoint)[0], 'results')
if not os.path.exists(save_path):
os.mkdir(save_path)
data_test = SDDData(width=CPI_WIDTH,
height=CPI_HEIGHT,
split='test',
normalize=False,
root=args.data_dir)
metrics = {"car": defaultdict(list), "ped": defaultdict(list)}
for scene_id in range(len(data_test.dataset.scenes)):
data_test.test_id = scene_id
print("scene", scene_id, "n_datas", len(data_test))
test_loader = torch.utils.data.DataLoader(dataset=data_test,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=True)
pedestrian_gt = []
car_gt = []
pedestrian_x = None
car_x = None
for bidx, data in enumerate(test_loader):
x, y_gt = data
if bidx % 2 == 0:
pedestrian_x = x
pedestrian_gt.append(y_gt)
else:
car_x = x
car_gt.append(y_gt)
for bidx, (obj_name, x, y_gt) in enumerate([("ped", pedestrian_x,
pedestrian_gt),
("car", car_x, car_gt)]):
y_gt = torch.stack(y_gt).float().to(args.gpu)
x = x.float().to(args.gpu)
_, y_pred = model.decode(x, 100)
log_py, log_px, _ = model.get_logprob(x.repeat(len(y_gt), 1, 1, 1),
y_gt)
log_py = log_py.cpu().detach().numpy().squeeze()
log_px = log_px.cpu().detach().numpy().squeeze()
metrics[obj_name]["nll_px"].extend(-1.0 * log_px)
metrics[obj_name]["nll_py"].extend(-1.0 * log_py)
y_gt_np = y_gt.detach().cpu().numpy().reshape((-1, 2))
y_pred = y_pred.cpu().detach().numpy().squeeze()
oracle_err = np.array([
mmfp_utils.compute_oracle_FDE(
y_pred.reshape(1, *y_pred.shape, 1, 1),
yg.reshape(
1,
1,
2,
1,
)) for yg in y_gt_np
])
metrics[obj_name]["oracle_err"].append(oracle_err.mean())
hist_gt, *_ = np.histogram2d(y_gt_np[:, 0],
y_gt_np[:, 1],
bins=np.linspace(0, 512, 512))
hist_pred, *_ = np.histogram2d(y_pred[:, 0],
y_pred[:, 1],
bins=np.linspace(0, 512, 512))
wemd = computeWEMD(hist_pred, hist_gt)
metrics[obj_name]["wemd"].append(wemd)
log_metrics = {
"oracle_err": oracle_err.mean(),
"wemd": wemd,
"nll_px": (-1 * log_px).mean(),
"nll_py": (-1 * log_py).mean()
}
print(f"scene {scene_id}", obj_name, log_metrics)
testing_sequence = data_test.dataset.scenes[
data_test.test_id].sequences[bidx]
objects_list = []
for k in range(3):
objects_list.append(
decode_obj(testing_sequence.objects[k],
testing_sequence.id))
objects = np.stack(objects_list, axis=0)
gt_object = np.array([[[[0], [0], [0], [0], [bidx]]]]).astype(
float) # mock it and draw dots instead
drawn_img_hyps = draw_hyps(testing_sequence.imgs[2],
y_pred,
gt_object,
objects,
normalize=False)
for (x1, y1) in y_gt_np:
color = (255, 0, 0)
cv2.circle(drawn_img_hyps, (x1, y1), 3, color, -1)
cv2.imwrite(
os.path.join(save_path,
f"{scene_id}-{bidx}-{obj_name}-hyps.jpg"),
drawn_img_hyps)
total_metrics = defaultdict(list)
for k, mets in metrics.items():
for obj_name, nums in mets.items():
print(f"Mean {k} {obj_name}: ", np.array(nums).mean())
total_metrics[obj_name].extend(nums)
for obj_name, nums in mets.items():
print(f"Total mean {obj_name}: ", np.array(nums).mean())
def main(args):
model = HyperRegression(args)
model = model.cuda()
resume_checkpoint = args.resume_checkpoint
print(f"Resume Path: {resume_checkpoint}")
checkpoint = torch.load(resume_checkpoint, map_location="cpu")
model_serialize = checkpoint['model']
model.load_state_dict(model_serialize)
model.eval()
save_path = os.path.join(os.path.split(resume_checkpoint)[0], 'results')
if not os.path.exists(save_path):
os.mkdir(save_path)
data_test = SDDData(split='test', normalize=False, root=args.data_dir)
nll_px_sum = 0
nll_py_sum = 0
multimod_emd_sum = 0
counter = 0.0
results = []
for session_id in range(len(data_test.dataset.scenes)):
data_test.test_id = session_id
test_loader = torch.utils.data.DataLoader(
dataset=data_test, batch_size=1, shuffle=False,
num_workers=0, pin_memory=True)
for bidx, data in enumerate(test_loader):
x, y_gt = data
x = x.float().to(args.gpu)
y_gt = y_gt.float().to(args.gpu).unsqueeze(1)
_, y_pred = model.decode(x, 1000)
log_py, log_px, _ = model.get_logprob(x, y_gt)
log_py = log_py.cpu().detach().numpy().squeeze()
log_px = log_px.cpu().detach().numpy().squeeze()
hyps_name = f"{session_id}-{bidx}-hyps.jpg"
print(hyps_name)
print("nll_x", str(-1.0 * log_px))
print("nll_y", str(-1.0 * log_py))
print("nll_(x+y)", str(-1.0 * log_px + log_py))
nll_px_sum = nll_px_sum + -1.0 * log_px
nll_py_sum = nll_py_sum + -1.0 * log_py
counter = counter + 1.0
y_pred = y_pred.cpu().detach().numpy().squeeze()
testing_sequence = data_test.dataset.scenes[data_test.test_id].sequences[bidx]
objects_list = []
for k in range(3):
objects_list.append(decode_obj(testing_sequence.objects[k], testing_sequence.id))
objects = np.stack(objects_list, axis=0)
gt_object = decode_obj(testing_sequence.objects[-1], testing_sequence.id)
drawn_img_hyps = draw_hyps(testing_sequence.imgs[-1], y_pred, gt_object, objects, normalize=False)
cv2.imwrite(os.path.join(save_path, hyps_name), drawn_img_hyps)
multimod_emd = mmfp_utils.wemd_from_pred_samples(y_pred)
multimod_emd_sum += multimod_emd
print("multimod_emd", multimod_emd)
_, _, height, width = x.shape
(X, Y), (log_px_grid, log_py_grid) = get_grid_logprob(height, width, x, model)
draw_sdd_heatmap(
objects=objects_list,
gt_object=gt_object,
testing_sequence=testing_sequence,
log_px_pred=log_px_grid,
X=X, Y=Y,
save_path=os.path.join(save_path, f"{session_id}-{bidx}-heatmap.png")
)
result_row = {
"session_id": session_id,
"bidx": bidx,
"nll_x": float(-1.0 * log_px),
"nll_y": float(-1.0 * log_py),
"multimod_emd": float(multimod_emd)
}
results.append(result_row)
print("Mean log_p_x: ", nll_px_sum / counter)
print("Mean log_p_y: ", nll_py_sum / counter)
print("Mean multimod_emd:", multimod_emd_sum / counter)
with open(os.path.join(save_path, "metrics.json"), "w") as f:
json.dump(results, f)
def main_worker(gpu, save_dir, ngpus_per_node, args):
# basic setup
cudnn.benchmark = True
normalize = False
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
model = HyperRegression(args)
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
start_epoch = 0
optimizer = model.make_optimizer(args)
if args.resume_checkpoint is None and os.path.exists(
os.path.join(save_dir, 'checkpoint-latest.pt')):
args.resume_checkpoint = os.path.join(
save_dir, 'checkpoint-latest.pt') # use the latest checkpoint
if args.resume_checkpoint is not None:
if args.resume_optimizer:
model, optimizer, start_epoch = resume(
args.resume_checkpoint,
model,
optimizer,
strict=(not args.resume_non_strict))
else:
model, _, start_epoch = resume(args.resume_checkpoint,
model,
optimizer=None,
strict=(not args.resume_non_strict))
print('Resumed from: ' + args.resume_checkpoint)
# initialize datasets and loaders
# initialize the learning rate scheduler
if args.scheduler == 'exponential':
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, args.exp_decay)
elif args.scheduler == 'step':
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=args.epochs // 2,
gamma=0.1)
elif args.scheduler == 'linear':
def lambda_rule(ep):
lr_l = 1.0 - max(0, ep - 0.5 * args.epochs) / float(
0.5 * args.epochs)
return lr_l
scheduler = optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda_rule)
else:
assert 0, "args.schedulers should be either 'exponential' or 'linear'"
# main training loop
start_time = time.time()
entropy_avg_meter = AverageValueMeter()
latent_nats_avg_meter = AverageValueMeter()
point_nats_avg_meter = AverageValueMeter()
if args.distributed:
print("[Rank %d] World size : %d" % (args.rank, dist.get_world_size()))
print("Start epoch: %d End epoch: %d" % (start_epoch, args.epochs))
data = SDDData(split='train', normalize=normalize, root=args.data_dir)
data_test = SDDData(split='test', normalize=normalize, root=args.data_dir)
train_loader = torch.utils.data.DataLoader(dataset=data,
batch_size=args.batch_size,
shuffle=True,
num_workers=0,
pin_memory=True)
test_loader = torch.utils.data.DataLoader(dataset=data_test,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=True)
for epoch in range(start_epoch, args.epochs):
# adjust the learning rate
if (epoch + 1) % args.exp_decay_freq == 0:
scheduler.step(epoch=epoch)
# train for one epoch
print("Epoch starts:")
for bidx, data in enumerate(train_loader):
# if bidx < 2:
x, y = data
#y = y.float().to(args.gpu).unsqueeze(1).repeat(1, 10).unsqueeze(2)
x = x.float().to(args.gpu)
y = y.float().to(args.gpu).unsqueeze(1)
y = y.repeat(1, 20, 1)
y += torch.randn(y.shape[0], y.shape[1], y.shape[2]).to(args.gpu)
step = bidx + len(train_loader) * epoch
model.train()
recon_nats = model(x, y, optimizer, step, None)
point_nats_avg_meter.update(recon_nats.item())
if step % args.log_freq == 0:
duration = time.time() - start_time
start_time = time.time()
print(
"[Rank %d] Epoch %d Batch [%2d/%2d] Time [%3.2fs] PointNats %2.5f"
% (args.rank, epoch, bidx, len(train_loader), duration,
point_nats_avg_meter.avg))
# print("Memory")
# print(process.memory_info().rss / (1024.0 ** 3))
# save visualizations
if (epoch + 1) % args.viz_freq == 0:
# reconstructions
model.eval()
for bidx, data in enumerate(test_loader):
x, _ = data
x = x.float().to(args.gpu)
_, y_pred = model.decode(x, 100)
y_pred = y_pred.cpu().detach().numpy().squeeze()
# y_pred[y_pred < 0] = 0
# y_pred[y_pred >= 0.98] = 0.98
testing_sequence = data_test.dataset.scenes[
data_test.test_id].sequences[bidx]
objects_list = []
for k in range(3):
objects_list.append(
decode_obj(testing_sequence.objects[k],
testing_sequence.id))
objects = np.stack(objects_list, axis=0)
gt_object = decode_obj(testing_sequence.objects[-1],
testing_sequence.id)
drawn_img_hyps = draw_hyps(testing_sequence.imgs[-1], y_pred,
gt_object, objects, normalize)
cv2.imwrite(
os.path.join(save_dir, 'images',
str(bidx) + '-' + str(epoch) + '-hyps.jpg'),
drawn_img_hyps)
if (epoch + 1) % args.save_freq == 0:
save(model, optimizer, epoch + 1,
os.path.join(save_dir, 'checkpoint-%d.pt' % epoch))
save(model, optimizer, epoch + 1,
os.path.join(save_dir, 'checkpoint-latest.pt'))