"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)
decoder = FullyConnectedTrajectoryTensorClassifier(
**decoder_args).to(DEVICE)
if multi_view_tensors:
multi_view = "multi_view"
else:
multi_view = "single_view"
embedder_path = os.path.join(
CROSS_VALIDATION_MODELS_PATH_AUTOENCODER,
multi_view,
"size_" + str(heatmap_size[0]),
"sigma_" + str(heatmap_smoothing_sigma),
"encoder_fold_" + str(fold) + ".weights",
)
"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"
embedder_path = os.path.join(
CROSS_VALIDATION_MODELS_PATH_AUTOENCODER,
multi_view,
"size_" + str(heatmap_size[0]),
"sigma_" + str(heatmap_smoothing_sigma),
"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 ########## #
embedder = CNN_2D_Encoder(**embedder_args).to(DEVICE)
encoder = TrajectoryTensorGRU(**encoder_args).to(DEVICE)
decoder = RecurrentDecoderTrajectoryTensor(**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")))
embedder_path = os.path.join(
CROSS_VALIDATION_MODELS_PATH_AUTOENCODER,
"multi_view",
"size_" + str(heatmap_size[0]),
"sigma_" + str(heatmap_smoothing_sigma),
"encoder_fold_" + str(fold) + ".weights",
)
embedder.load_state_dict(torch.load(embedder_path))
params = list(encoder.parameters()) + list(decoder.parameters()) + list(embedder.parameters())