"num_workers": NUM_WORKERS,
}
test_dataset_args = {
"inputs_path": trajectory_tensor_path,
"targets_path": CROSS_VALIDATION_WHERE_TARGETS_PATH,
"fold": fold,
"heatmap_size": heatmap_size,
"heatmap_smoothing_sigma": heatmap_smoothing_sigma,
"phase": "test",
"batch_size": BATCH_SIZE,
"shuffle": False,
"num_workers": NUM_WORKERS,
}
train_loader = get_trajectory_tensor_dataset(**train_dataset_args)
val_loader = get_trajectory_tensor_dataset(**val_dataset_args)
test_loader = get_trajectory_tensor_dataset(**test_dataset_args)
# ########## SET UP MODEL ########## #
embedder = CNN_2D_Encoder(**embedder_args).to(DEVICE)
encoder = TrajectoryTensorGRU(**encoder_args).to(DEVICE)
temporal_decoder = RecurrentDecoderTrajectoryTensor(
**temporal_decoder_args).to(DEVICE)
spatial_decoder = Trajectory_Tensor_CNN_2D_Decoder(
**spatial_decoder_args).to(DEVICE)
if multi_view_tensors:
multi_view = "multi_view"
else:
multi_view = "single_view"
# ########## SET UP DATASET ########## #
test_dataset_args = {
"inputs_path": trajectory_tensor_path,
"targets_path": when_targets_path,
"fold": fold,
"heatmap_size": heatmap_size,
"heatmap_smoothing_sigma": heatmap_smoothing_sigma,
"phase": "test",
"batch_size": BATCH_SIZE,
"shuffle": False,
"num_workers": NUM_WORKERS,
"multi_target": True,
}
test_loader = get_trajectory_tensor_dataset(**test_dataset_args)
# ########## SET UP MODEL ########## #
encoder = CNN_2D_1D_Encoder(**encoder_args).to(DEVICE)
decoder = CNN_1D_Trajectory_Tensor_Classifier(**decoder_args).to(DEVICE)
encoder.load_state_dict(torch.load(os.path.join(MODEL_LOAD_PATH, "encoder_fold_" + str(fold) + ".weights")))
decoder.load_state_dict(torch.load(os.path.join(MODEL_LOAD_PATH, "decoder_fold_" + str(fold) + ".weights")))
params = list(encoder.parameters()) + list(decoder.parameters())
loss_function = nn.BCELoss()
# ########## TRAIN AND EVALUATE ########## #
best_ap = 0
test_args = {
