ch.setLevel(logging.WARN)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
ch.setFormatter(formatter)
logger.addHandler(fh)
logger.addHandler(ch)
logger.info(args)
### Extract data for training, validation and testing
node_features, edge_features, full_data, train_data, val_data, test_data, new_node_val_data, \
new_node_test_data = get_data(DATA,
different_new_nodes_between_val_and_test=args.different_new_nodes)
# Initialize training neighbor finder to retrieve temporal graph
train_ngh_finder = get_neighbor_finder(train_data, args.uniform)
# Initialize validation and test neighbor finder to retrieve temporal graph
full_ngh_finder = get_neighbor_finder(full_data, args.uniform)
# Initialize negative samplers. Set seeds for validation and testing so negatives are the same
# across different runs
# NB: in the inductive setting, negatives are sampled only amongst other new nodes
train_rand_sampler = RandEdgeSampler(train_data.sources,
train_data.destinations)
val_rand_sampler = RandEdgeSampler(full_data.sources,
full_data.destinations,
seed=0)
nn_val_rand_sampler = RandEdgeSampler(new_node_val_data.sources,
new_node_val_data.destinations,
seed=1)
ch.setLevel(logging.WARN)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh.setFormatter(formatter)
ch.setFormatter(formatter)
logger.addHandler(fh)
logger.addHandler(ch)
logger.info(args)
full_data, node_features, edge_features, train_data, val_data, test_data = \
get_data_node_classification(DATA, use_validation=args.use_validation, vector_features=args.vector_features)
max_idx = max(full_data.unique_nodes)
train_ngh_finder = get_neighbor_finder(train_data,
uniform=UNIFORM,
max_node_idx=max_idx)
# Set device
device_string = 'cuda:{}'.format(GPU) if torch.cuda.is_available() else 'cpu'
device = torch.device(device_string)
# Compute time statistics
mean_time_shift_src, std_time_shift_src, mean_time_shift_dst, std_time_shift_dst = \
compute_time_statistics(full_data.sources, full_data.destinations, full_data.timestamps)
# class
class Net(nn.Module):
def __init__(self, neighbor_finder, node_features, edge_features, device,
n_layers, n_heads, dropout, use_memory, message_dimension,