def main():
"""Load data and perform experiments."""
args = parse_arguments()
if not baseline_utils.validate_args(args):
return
np.random.seed(100)
# Get features
if args.use_map and args.use_social:
baseline_key = "map_social"
elif args.use_map:
baseline_key = "map"
elif args.use_social:
baseline_key = "social"
else:
baseline_key = "none"
# Get data
data_dict = baseline_utils.get_data(args, baseline_key)
# Perform experiments
start = time.time()
perform_k_nn_experiments(data_dict, baseline_key)
end = time.time()
print(f"Completed experiment in {(end-start)/60.0} mins")
def main():
"""Main."""
args = parse_arguments()
if not baseline_utils.validate_args(args):
return
print(f"Using all ({joblib.cpu_count()}) CPUs....")
if use_cuda:
print(f"Using all ({torch.cuda.device_count()}) GPUs...")
model_utils = ModelUtils()
# key for getting feature set
# Get features
if args.use_map and args.use_social:
baseline_key = "map_social"
elif args.use_map:
baseline_key = "map"
elif args.use_social:
baseline_key = "social"
else:
baseline_key = "none"
# Get data
data_dict = baseline_utils.get_data(args, baseline_key)
# Get model
criterion = nn.MSELoss()
encoder = EncoderRNN(
input_size=len(baseline_utils.BASELINE_INPUT_FEATURES[baseline_key]))
decoder = DecoderRNN(output_size=2)
if use_cuda:
encoder = nn.DataParallel(encoder)
decoder = nn.DataParallel(decoder)
encoder.to(device)
decoder.to(device)
encoder_optimizer = torch.optim.Adam(encoder.parameters(), lr=args.lr)
decoder_optimizer = torch.optim.Adam(decoder.parameters(), lr=args.lr)
# If model_path provided, resume from saved checkpoint
if args.model_path is not None and os.path.isfile(args.model_path):
epoch, rollout_len, _ = model_utils.load_checkpoint(
args.model_path, encoder, decoder, encoder_optimizer,
decoder_optimizer)
start_epoch = epoch + 1
start_rollout_idx = ROLLOUT_LENS.index(rollout_len) + 1
else:
start_epoch = 0
start_rollout_idx = 0
if not args.test:
# Tensorboard logger
log_dir = os.path.join(os.getcwd(), "lstm_logs", baseline_key)
# Get PyTorch Dataset
train_dataset = LSTMDataset(data_dict, args, "train")
val_dataset = LSTMDataset(data_dict, args, "val")
# Setting Dataloaders
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.train_batch_size,
shuffle=True,
drop_last=False,
collate_fn=model_utils.my_collate_fn,
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.val_batch_size,
drop_last=False,
shuffle=False,
collate_fn=model_utils.my_collate_fn,
)
print("Training begins ...")
decrement_counter = 0
epoch = start_epoch
global_start_time = time.time()
for i in range(start_rollout_idx, len(ROLLOUT_LENS)):
rollout_len = ROLLOUT_LENS[i]
logger = Logger(log_dir, name="{}".format(rollout_len))
best_loss = float("inf")
prev_loss = best_loss
while epoch < args.end_epoch:
start = time.time()
train(
train_loader,
epoch,
criterion,
logger,
encoder,
decoder,
encoder_optimizer,
decoder_optimizer,
model_utils,
rollout_len,
)
end = time.time()
print(
f"Training epoch completed in {(end - start) / 60.0} mins, Total time: {(end - global_start_time) / 60.0} mins"
)
epoch += 1
if epoch % 5 == 0:
start = time.time()
prev_loss, decrement_counter = validate(
val_loader,
epoch,
criterion,
logger,
encoder,
decoder,
encoder_optimizer,
decoder_optimizer,
model_utils,
prev_loss,
decrement_counter,
rollout_len,
)
end = time.time()
print(
f"Validation completed in {(end - start) / 60.0} mins, Total time: {(end - global_start_time) / 60.0} mins"
)
# If val loss increased 3 times consecutively, go to next rollout length
if decrement_counter > 2:
break
else:
start_time = time.time()
temp_save_dir = tempfile.mkdtemp()
test_size = data_dict["test_input"].shape[0]
test_data_subsets = baseline_utils.get_test_data_dict_subset(
data_dict, args)
# test_batch_size should be lesser than joblib_batch_size
Parallel(n_jobs=1, verbose=2)(
delayed(infer_helper)(test_data_subsets[i], i, encoder, decoder,
model_utils, temp_save_dir)
for i in range(0, test_size, args.joblib_batch_size))
baseline_utils.merge_saved_traj(temp_save_dir, args.traj_save_path)
shutil.rmtree(temp_save_dir)
end = time.time()
print(f"Test completed in {(end - start_time) / 60.0} mins")
print(f"Forecasted Trajectories saved at {args.traj_save_path}")