Path("results/").mkdir(parents=True, exist_ok=True)
# Initialize Model
tgn = TGN(
neighbor_finder=train_ngh_finder,
node_features=node_features,
edge_features=edge_features,
device=device,
n_layers=NUM_LAYER,
n_heads=NUM_HEADS,
dropout=DROP_OUT,
use_memory=USE_MEMORY,
message_dimension=MESSAGE_DIM,
memory_dimension=MEMORY_DIM,
memory_update_at_start=not args.memory_update_at_end,
embedding_module_type=args.embedding_module,
message_function=args.message_function,
aggregator_type=args.aggregator,
memory_updater_type=args.memory_updater,
n_neighbors=NUM_NEIGHBORS,
mean_time_shift_src=mean_time_shift_src,
std_time_shift_src=std_time_shift_src,
mean_time_shift_dst=mean_time_shift_dst,
std_time_shift_dst=std_time_shift_dst,
use_destination_embedding_in_message=args.
use_destination_embedding_in_message,
use_source_embedding_in_message=args.use_source_embedding_in_message,
dyrep=args.dyrep)
# loss function
criterion = torch.nn.BCELoss()
args.prefix)
Path("results/").mkdir(parents=True, exist_ok=True)
# Initialize Model
tgn = TGN(
neighbor_finder=train_ngh_finder,
node_features=node_features,
edge_features=edge_features,
device=device,
n_layers=NUM_LAYER,
n_heads=NUM_HEADS,
dropout=DROP_OUT,
use_memory=USE_MEMORY,
message_dimension=MESSAGE_DIM,
memory_dimension=MEMORY_DIM,
memory_update_at_start=not args.memory_update_at_end,
embedding_module_type=args.embedding_module,
message_function=args.message_function,
aggregator_type=args.aggregator,
n_neighbors=NUM_NEIGHBORS,
mean_time_shift_src=mean_time_shift_src,
std_time_shift_src=std_time_shift_src,
mean_time_shift_dst=mean_time_shift_dst,
std_time_shift_dst=std_time_shift_dst,
use_destination_embedding_in_message=args.
use_destination_embedding_in_message,
use_source_embedding_in_message=args.use_source_embedding_in_message)
tgn = tgn.to(device)
num_instance = len(train_data.sources)