valid_test_new_old_flag], dst_l[valid_test_new_old_flag], ts_l[
valid_test_new_old_flag], e_idx_l[
valid_test_new_old_flag], label_l[valid_test_new_old_flag]
train_data = train_src_l, train_dst_l, train_ts_l, train_e_idx_l, train_label_l
val_data = val_src_l, val_dst_l, val_ts_l, val_e_idx_l, val_label_l
train_val_data = (train_data, val_data)
# create two neighbor finders to handle graph extraction.
# for transductive mode all phases use full_ngh_finder, for inductive node train/val phases use the partial one
# while test phase still always uses the full one
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,
bias=args.bias,
use_cache=NGH_CACHE,
sample_method=args.pos_sample)
partial_adj_list = [[] for _ in range(max_idx + 1)]
for src, dst, eidx, ts in zip(train_src_l, train_dst_l, train_e_idx_l,
train_ts_l):
partial_adj_list[src].append((dst, eidx, ts))
partial_adj_list[dst].append((src, eidx, ts))
for src, dst, eidx, ts in zip(val_src_l, val_dst_l, val_e_idx_l, val_ts_l):
partial_adj_list[src].append((dst, eidx, ts))
partial_adj_list[dst].append((src, eidx, ts))
partial_ngh_finder = NeighborFinder(partial_adj_list,
bias=args.bias,
use_cache=NGH_CACHE,
sample_method=args.pos_sample)
ngh_finders = partial_ngh_finder, full_ngh_finder
val_ts_l = val_data.ts.values
val_label_l = val_data.label.values
test_src_l = test_data.u.values
test_dst_l = test_data.i.values
test_ts_l = test_data.ts.values
test_label_l = test_data.label.values
# Initialize the data structure for graph and edge sampling
# build the graph for fast query
adj_list = [[] for _ in range(max_idx + 1)]
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)
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))
else:
mult = 2
return Classifier(args,
in_features=args.gcn_out_feats * mult,
out_features=2).cuda()
if __name__ == '__main__':
parser = u.create_parser()
args = u.parse_args(parser)
args.device = 'cuda'
dataset = build_dataset(args)
tasker = build_tasker(args, dataset)
splitter = splitter(args, tasker)
adj_list = u.build_adj_list(dataset)
ngh_finder = NeighborFinder(adj_list)
gcn = TGAN(ngh_finder, args.gcn_out_feats, dataset.num_nodes,
args.num_hist_steps)
classifier = build_classifier(args)
loss = ce.Cross_Entropy(args, dataset).to('cuda')
trainer = Trainer(args,
splitter=splitter,
gcn=gcn,
classifier=classifier,
comp_loss=loss,
dataset=dataset,
num_classes=dataset.num_classes)
trainer.train()
