"targets_path": which_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 = FullyConnectedTrajectoryTensorClassifier(**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 = {
"encoder": encoder,
"decoder": decoder,
"device": DEVICE,
"test_loader": test_loader,
"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 ########## #
encoder = CNN_3D_Encoder(**encoder_args).to(DEVICE)
decoder = FullyConnectedTrajectoryTensorClassifier(
**decoder_args).to(DEVICE)
params = list(encoder.parameters()) + list(decoder.parameters())
optimizer = optim.Adam(params,
lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY)
loss_function = nn.BCELoss()
# ########## TRAIN AND EVALUATE ########## #
best_ap = 0
for epoch in range(NUM_EPOCHS):
print("----------- EPOCH " + str(epoch) + " -----------")
trainer_args = {
"encoder": encoder,