full_ngh_finder = NeighborFinder(full_adj_list, uniform=UNIFORM)
train_rand_sampler = RandEdgeSampler(train_src_l, train_dst_l)
val_rand_sampler = RandEdgeSampler(src_l, dst_l)
test_rand_sampler = RandEdgeSampler(src_l, dst_l)
# Model initialize
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device('cuda:{}'.format(GPU))
tgan = TGAN(train_ngh_finder,
n_feat,
e_feat,
n_feat_freeze=args.freeze,
num_layers=NUM_LAYER,
use_time=USE_TIME,
agg_method=AGG_METHOD,
attn_mode=ATTN_MODE,
seq_len=SEQ_LEN,
n_head=NUM_HEADS,
drop_out=DROP_OUT,
node_dim=NODE_DIM,
time_dim=TIME_DIM)
optimizer = torch.optim.Adam(tgan.parameters(), lr=LEARNING_RATE)
criterion = torch.nn.BCELoss()
tgan = tgan.to(device)
num_instance = len(train_src_l)
num_batch = math.ceil(num_instance / BATCH_SIZE)
logger.info('num of training instances: {}'.format(num_instance))
logger.info('num of batches per epoch: {}'.format(num_batch))
#for src, dst, eidx, ts in zip(train_src_l, train_dst_l, train_e_idx_l, train_ts_l):
# adj_list[src].append((dst, eidx, ts))
# adj_list[dst].append((src, eidx, ts))
#train_ngh_finder = NeighborFinder(adj_list, uniform=UNIFORM)
# full graph with all the data for the test and validation purpose
full_adj_list = [[] for _ in range(max_idx + 1)]
for src, dst, eidx, ts in zip(src_l, dst_l, e_idx_l, ts_l):
full_adj_list[src].append((dst, eidx, ts))
full_adj_list[dst].append((src, eidx, ts))
full_ngh_finder = NeighborFinder(full_adj_list, uniform=UNIFORM)
### Model initialize
device = torch.device('cuda:{}'.format(GPU))
tgan = TGAN(full_ngh_finder, n_feat, e_feat,
num_layers=NUM_LAYER, use_time=USE_TIME, agg_method=AGG_METHOD, attn_mode=ATTN_MODE,
seq_len=SEQ_LEN, n_head=NUM_HEADS, drop_out=DROP_OUT, node_dim=NODE_DIM, time_dim=TIME_DIM)
# optimizer = torch.optim.Adam(tgan.parameters(), lr=LEARNING_RATE)
# criterion = torch.nn.BCELoss()
tgan = tgan.to(device)
#logger.info('loading saved TGAN model')
model_path = f'./saved_models/{args.prefix}-{args.agg_method}-{args.attn_mode}-{DATA}.pth'
tgan.load_state_dict(torch.load(model_path))
tgan.eval()
#logger.info('TGAN models loaded')
#logger.info('TGAN starts generating representations')
with torch.no_grad():
tgan.eval()
src_embed = tgan.tem_conv(src_l, ts_l, NODE_LAYER)