def main():
args = parse_arguments()
model_utils = ModelUtils()
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"
data_dict = baseline_utils.get_data(args, baseline_key)
criterion = nn.MSELoss()
encoder = EncoderRNN(ins=len(baseline_utils.BASELINE_INPUT_FEATURES[baseline_key]))
decoder = DecoderRNN(ops=2)
if use_cuda:
encoder = nn.DataParallel(encoder)
decoder = nn.DataParallel(decoder)
encoder.to(device)
decoder.to(device)
encoder_optim = torch.optim.Adam(encoder.parameters(), lr=args.lr)
decoder_optim = torch.optim.Adam(decoder.parameters(), lr=args.lr)
rollout_id = 0
log_dir = os.path.join(os.getcwd(), "lstm_logs", baseline_key)
train_dataset = LSTMDataset(data_dict, args, "train")
val_dataset = LSTMDataset(data_dict, args, "val")
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=True, collate_fn=model_utils.my_collate_fn,)
epoch = 0
for i in range(rollout_id, len(ROLLOUT_LENS)):
rollout_len = ROLLOUT_LENS[i]
logger = Logger(log_dir, name="{}".format(rollout_len))
while epoch < args.end_epoch:
start = time.time()
train(train_loader, epoch, criterion, logger, encoder, decoder, encoder_optim, decoder_optim, model_utils, rollout_len)
epoch += 1
if epoch % 5 == 0:
val_loss = validate(val_loader, epoch, criterion, logger, encoder, decoder, encoder_optim, decoder_optim, model_utils, rollout_len)
print("Validation loss =", val_loss, "after", epoch, "epochs")
def infer_helper(
curr_data_dict: Dict[str, Any],
start_idx: int,
encoder: EncoderRNN,
decoder: DecoderRNN,
model_utils: ModelUtils,
forecasted_save_dir: str,
):
"""Run inference on the current joblib batch.
Args:
curr_data_dict: Data dictionary for the current joblib batch
start_idx: Start idx of the current joblib batch
encoder: Encoder network instance
decoder: Decoder network instance
model_utils: ModelUtils instance
forecasted_save_dir: Directory where forecasted trajectories are to be saved
"""
args = parse_arguments()
curr_test_dataset = LSTMDataset(curr_data_dict, args, "test")
curr_test_loader = torch.utils.data.DataLoader(
curr_test_dataset,
shuffle=False,
batch_size=args.test_batch_size,
collate_fn=model_utils.my_collate_fn,
)
if args.use_map:
print(f"#### LSTM+map inference at index {start_idx} ####")
infer_map(
curr_test_loader,
encoder,
decoder,
start_idx,
forecasted_save_dir,
model_utils,
)
else:
print(f"#### LSTM+social inference at {start_idx} ####"
) if args.use_social else print(
f"#### LSTM inference at {start_idx} ####")
infer_absolute(
curr_test_loader,
encoder,
decoder,
start_idx,
forecasted_save_dir,
model_utils,
)
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}")