def infer_map(
test_loader: torch.utils.data.DataLoader,
encoder: EncoderRNN,
decoder: DecoderRNN,
start_idx: int,
forecasted_save_dir: str,
model_utils: ModelUtils,
):
"""Infer function for map-based LSTM baselines and save the forecasted trajectories.
Args:
test_loader: DataLoader for the test set
encoder: Encoder network instance
decoder: Decoder network instance
start_idx: start index for the current joblib batch
forecasted_save_dir: Directory where forecasted trajectories are to be saved
model_utils: ModelUtils instance
"""
args = parse_arguments()
global best_loss
forecasted_trajectories = {}
for i, (_input, target, helpers) in enumerate(test_loader):
_input = _input.to(device)
batch_helpers = list(zip(*helpers))
helpers_dict = {}
for k, v in config.LSTM_HELPER_DICT_IDX.items():
helpers_dict[k] = batch_helpers[v]
# Set to eval mode
encoder.eval()
decoder.eval()
# Encoder
batch_size = _input.shape[0]
input_length = _input.shape[1]
# Iterate over every element in the batch
for batch_idx in range(batch_size):
num_candidates = len(
helpers_dict["CANDIDATE_CENTERLINES"][batch_idx])
curr_centroids = helpers_dict["CENTROIDS"][batch_idx]
seq_id = int(helpers_dict["SEQ_PATHS"][batch_idx])
abs_outputs = []
# Predict using every centerline candidate for the current trajectory
for candidate_idx in range(num_candidates):
curr_centerline = helpers_dict["CANDIDATE_CENTERLINES"][
batch_idx][candidate_idx]
curr_nt_dist = helpers_dict["CANDIDATE_NT_DISTANCES"][
batch_idx][candidate_idx]
_input = torch.FloatTensor(
np.expand_dims(curr_nt_dist[:args.obs_len].astype(float),
0)).to(device)
# Initialize encoder hidden state
encoder_hidden = model_utils.init_hidden(
1, encoder.module.hidden_size
if use_cuda else encoder.hidden_size)
# Encode observed trajectory
for ei in range(input_length):
encoder_input = _input[:, ei, :]
encoder_hidden = encoder(encoder_input, encoder_hidden)
# Initialize decoder input with last coordinate in encoder
decoder_input = encoder_input[:, :2]
# Initialize decoder hidden state as encoder hidden state
decoder_hidden = encoder_hidden
decoder_outputs = torch.zeros((1, args.pred_len, 2)).to(device)
# Decode hidden state in future trajectory
for di in range(args.pred_len):
decoder_output, decoder_hidden = decoder(
decoder_input, decoder_hidden)
decoder_outputs[:, di, :] = decoder_output
# Use own predictions as inputs at next step
decoder_input = decoder_output
# Get absolute trajectory
abs_helpers = {}
abs_helpers["REFERENCE"] = np.expand_dims(
np.array(helpers_dict["CANDIDATE_DELTA_REFERENCES"]
[batch_idx][candidate_idx]),
0,
)
abs_helpers["CENTERLINE"] = np.expand_dims(curr_centerline, 0)
abs_input, abs_output = baseline_utils.get_abs_traj(
_input.clone().cpu().numpy(),
decoder_outputs.detach().clone().cpu().numpy(),
args,
abs_helpers,
)
# array of shape (1,30,2) to list of (30,2)
abs_outputs.append(abs_output[0])
forecasted_trajectories[seq_id] = abs_outputs
os.makedirs(forecasted_save_dir, exist_ok=True)
with open(os.path.join(forecasted_save_dir, f"{start_idx}.pkl"),
"wb") as f:
pkl.dump(forecasted_trajectories, f)
def validate(
val_loader: Any,
epoch: int,
criterion: Any,
logger: Logger,
encoder: Any,
decoder: Any,
encoder_optimizer: Any,
decoder_optimizer: Any,
model_utils: ModelUtils,
prev_loss: float,
decrement_counter: int,
rollout_len: int = 30,
) -> Tuple[float, int]:
"""Validate the lstm network.
Args:
val_loader: DataLoader for the train set
epoch: epoch number
criterion: Loss criterion
logger: Tensorboard logger
encoder: Encoder network instance
decoder: Decoder network instance
encoder_optimizer: optimizer for the encoder network
decoder_optimizer: optimizer for the decoder network
model_utils: instance for ModelUtils class
prev_loss: Loss in the previous validation run
decrement_counter: keeping track of the number of consecutive times loss increased in the current rollout
rollout_len: current prediction horizon
"""
args = parse_arguments()
global best_loss
total_loss = []
for i, (_input, target, helpers) in enumerate(val_loader):
_input = _input.to(device)
target = target.to(device)
# Set to eval mode
encoder.eval()
decoder.eval()
# Encoder
batch_size = _input.shape[0]
input_length = _input.shape[1]
output_length = target.shape[1]
input_shape = _input.shape[2]
# Initialize encoder hidden state
encoder_hidden = model_utils.init_hidden(
batch_size,
encoder.module.hidden_size if use_cuda else encoder.hidden_size)
# Initialize loss
loss = 0
# Encode observed trajectory
for ei in range(input_length):
encoder_input = _input[:, ei, :]
encoder_hidden = encoder(encoder_input, encoder_hidden)
# Initialize decoder input with last coordinate in encoder
decoder_input = encoder_input[:, :2]
# Initialize decoder hidden state as encoder hidden state
decoder_hidden = encoder_hidden
decoder_outputs = torch.zeros(target.shape).to(device)
# Decode hidden state in future trajectory
for di in range(output_length):
decoder_output, decoder_hidden = decoder(decoder_input,
decoder_hidden)
decoder_outputs[:, di, :] = decoder_output
# Update losses for all benchmarks
loss += criterion(decoder_output[:, :2], target[:, di, :2])
# Use own predictions as inputs at next step
decoder_input = decoder_output
# Get average loss for pred_len
loss = loss / output_length
total_loss.append(loss)
if i % 10 == 0:
cprint(
f"Val -- Epoch:{epoch}, loss:{loss}, Rollout: {rollout_len}",
color="green",
)
# Save
val_loss = sum(total_loss) / len(total_loss)
if val_loss <= best_loss:
best_loss = val_loss
if args.use_map:
save_dir = "saved_models/lstm_map"
elif args.use_social:
save_dir = "saved_models/lstm_social"
else:
save_dir = "saved_models/lstm"
os.makedirs(save_dir, exist_ok=True)
model_utils.save_checkpoint(
save_dir,
{
"epoch": epoch + 1,
"rollout_len": rollout_len,
"encoder_state_dict": encoder.state_dict(),
"decoder_state_dict": decoder.state_dict(),
"best_loss": val_loss,
"encoder_optimizer": encoder_optimizer.state_dict(),
"decoder_optimizer": decoder_optimizer.state_dict(),
},
)
logger.scalar_summary(tag="Val/loss", value=val_loss.item(), step=epoch)
# Keep track of the loss to change preiction horizon
if val_loss <= prev_loss:
decrement_counter = 0
else:
decrement_counter += 1
return val_loss, decrement_counter
def infer_absolute(
test_loader: torch.utils.data.DataLoader,
encoder: EncoderRNN,
decoder: DecoderRNN,
start_idx: int,
forecasted_save_dir: str,
model_utils: ModelUtils,
):
"""Infer function for non-map LSTM baselines and save the forecasted trajectories.
Args:
test_loader: DataLoader for the test set
encoder: Encoder network instance
decoder: Decoder network instance
start_idx: start index for the current joblib batch
forecasted_save_dir: Directory where forecasted trajectories are to be saved
model_utils: ModelUtils instance
"""
args = parse_arguments()
forecasted_trajectories = {}
for i, (_input, target, helpers) in enumerate(test_loader):
_input = _input.to(device)
batch_helpers = list(zip(*helpers))
helpers_dict = {}
for k, v in config.LSTM_HELPER_DICT_IDX.items():
helpers_dict[k] = batch_helpers[v]
# Set to eval mode
encoder.eval()
decoder.eval()
# Encoder
batch_size = _input.shape[0]
input_length = _input.shape[1]
input_shape = _input.shape[2]
# Initialize encoder hidden state
encoder_hidden = model_utils.init_hidden(
batch_size,
encoder.module.hidden_size if use_cuda else encoder.hidden_size)
# Encode observed trajectory
for ei in range(input_length):
encoder_input = _input[:, ei, :]
encoder_hidden = encoder(encoder_input, encoder_hidden)
# Initialize decoder input with last coordinate in encoder
decoder_input = encoder_input[:, :2]
# Initialize decoder hidden state as encoder hidden state
decoder_hidden = encoder_hidden
decoder_outputs = torch.zeros(
(batch_size, args.pred_len, 2)).to(device)
# Decode hidden state in future trajectory
for di in range(args.pred_len):
decoder_output, decoder_hidden = decoder(decoder_input,
decoder_hidden)
decoder_outputs[:, di, :] = decoder_output
# Use own predictions as inputs at next step
decoder_input = decoder_output
# Get absolute trajectory
abs_helpers = {}
abs_helpers["REFERENCE"] = np.array(helpers_dict["DELTA_REFERENCE"])
abs_helpers["TRANSLATION"] = np.array(helpers_dict["TRANSLATION"])
abs_helpers["ROTATION"] = np.array(helpers_dict["ROTATION"])
abs_inputs, abs_outputs = baseline_utils.get_abs_traj(
_input.clone().cpu().numpy(),
decoder_outputs.detach().clone().cpu().numpy(),
args,
abs_helpers,
)
for i in range(abs_outputs.shape[0]):
seq_id = int(helpers_dict["SEQ_PATHS"][i])
forecasted_trajectories[seq_id] = [abs_outputs[i]]
with open(os.path.join(forecasted_save_dir, f"{start_idx}.pkl"),
"wb") as f:
pkl.dump(forecasted_trajectories, f)
def train(
train_loader: Any,
epoch: int,
criterion: Any,
logger: Logger,
encoder: Any,
decoder: Any,
encoder_optimizer: Any,
decoder_optimizer: Any,
model_utils: ModelUtils,
rollout_len: int = 30,
) -> None:
"""Train the lstm network.
Args:
train_loader: DataLoader for the train set
epoch: epoch number
criterion: Loss criterion
logger: Tensorboard logger
encoder: Encoder network instance
decoder: Decoder network instance
encoder_optimizer: optimizer for the encoder network
decoder_optimizer: optimizer for the decoder network
model_utils: instance for ModelUtils class
rollout_len: current prediction horizon
"""
args = parse_arguments()
global global_step
for i, (_input, target, helpers) in enumerate(train_loader):
_input = _input.to(device)
target = target.to(device)
# Set to train mode
encoder.train()
decoder.train()
# Zero the gradients
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
# Encoder
batch_size = _input.shape[0]
input_length = _input.shape[1]
output_length = target.shape[1]
input_shape = _input.shape[2]
# Initialize encoder hidden state
encoder_hidden = model_utils.init_hidden(
batch_size,
encoder.module.hidden_size if use_cuda else encoder.hidden_size)
# Initialize losses
loss = 0
# Encode observed trajectory
for ei in range(input_length):
encoder_input = _input[:, ei, :]
encoder_hidden = encoder(encoder_input, encoder_hidden)
# Initialize decoder input with last coordinate in encoder
decoder_input = encoder_input[:, :2]
# Initialize decoder hidden state as encoder hidden state
decoder_hidden = encoder_hidden
decoder_outputs = torch.zeros(target.shape).to(device)
# Decode hidden state in future trajectory
for di in range(rollout_len):
decoder_output, decoder_hidden = decoder(decoder_input,
decoder_hidden)
decoder_outputs[:, di, :] = decoder_output
# Update loss
loss += criterion(decoder_output[:, :2], target[:, di, :2])
# Use own predictions as inputs at next step
decoder_input = decoder_output
# Get average loss for pred_len
loss = loss / rollout_len
# Backpropagate
loss.backward()
encoder_optimizer.step()
decoder_optimizer.step()
if global_step % 1000 == 0:
# Log results
print(
f"Train -- Epoch:{epoch}, loss:{loss}, Rollout:{rollout_len}")
logger.scalar_summary(tag="Train/loss",
value=loss.item(),
step=epoch)
global_step += 1
