|
|
@@ -1,2 +1,77 @@ |
|
|
|
def trainprep(data):
|
|
|
|
pass
|
|
|
|
from .sampling import fixed_unigram_candidate_sampler
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
|
|
def train_val_test_split_edges(edges, ratios):
|
|
|
|
train_ratio, val_ratio, test_ratio = ratios
|
|
|
|
|
|
|
|
if train_ratio + val_ratio + test_ratio != 1.0:
|
|
|
|
raise ValueError('Train, validation and test ratios must add up to 1')
|
|
|
|
|
|
|
|
order = torch.randperm(len(edges))
|
|
|
|
edges = edges[order, :]
|
|
|
|
n = round(len(edges) * train_ratio)
|
|
|
|
edges_train = edges[:n]
|
|
|
|
n_1 = round(len(edges) * (train_ratio + val_ratio))
|
|
|
|
edges_val = edges[n:n_1]
|
|
|
|
edges_test = edges[n_1:]
|
|
|
|
|
|
|
|
return edges_train, edges_val, edges_test
|
|
|
|
|
|
|
|
|
|
|
|
def prepare_adj_mat(adj_mat, ratios):
|
|
|
|
degrees = adj_mat.sum(0)
|
|
|
|
edges_pos = torch.nonzero(adj_mat)
|
|
|
|
|
|
|
|
neg_neighbors = fixed_unigram_candidate_sampler(edges_pos[:, 1],
|
|
|
|
len(edges), degrees, 0.75)
|
|
|
|
edges_neg = torch.cat((edges_pos[:, 0], neg_neighbors.view(-1, 1)), 1)
|
|
|
|
|
|
|
|
edges_pos = (edges_pos_train, edges_pos_val, edges_pos_test) = \
|
|
|
|
train_val_test_split_edges(edges_pos, ratios)
|
|
|
|
edges_neg = (edges_neg_train, edges_neg_val, edges_neg_test) = \
|
|
|
|
train_val_test_split_edges(edges_neg, ratios)
|
|
|
|
|
|
|
|
return edges_pos, edges_neg
|
|
|
|
|
|
|
|
|
|
|
|
class PreparedRelation(object):
|
|
|
|
def __init__(self, node_type_row, node_type_column,
|
|
|
|
adj_mat_train, adj_mat_train_trans,
|
|
|
|
edges_pos, edges_neg, edges_pos_trans, edges_neg_trans):
|
|
|
|
|
|
|
|
self.adj_mat_train = adj_mat_train
|
|
|
|
self.adj_mat_train_trans = adj_mat_train_trans
|
|
|
|
self.edges_pos = edges_pos
|
|
|
|
self.edges_neg = edges_neg
|
|
|
|
self.edges_pos_trans = edges_pos_trans
|
|
|
|
self.edges_neg_trans = edges_neg_trans
|
|
|
|
|
|
|
|
|
|
|
|
def prepare_relation(r, ratios):
|
|
|
|
adj_mat = r.get_adjacency_matrix(r.node_type_row, r.node_type_column)
|
|
|
|
edges_pos, edges_neg = prepare_adj_mat(adj_mat)
|
|
|
|
|
|
|
|
# adj_mat_train = torch.zeros_like(adj_mat)
|
|
|
|
# adj_mat_train[edges_pos[0][:, 0], edges_pos[0][:, 0]] = 1
|
|
|
|
adj_mat_train = torch.sparse_coo_tensor(indices = edges_pos[0].transpose(0, 1),
|
|
|
|
values=torch.ones(len(edges_pos[0]), dtype=adj_mat.dtype))
|
|
|
|
|
|
|
|
if r.node_type_row != r.node_type_col:
|
|
|
|
adj_mat_trans = r.get_adjacency_matrix(r.node_type_col, r.node_type_row)
|
|
|
|
edges_pos_trans, edges_neg_trans = prepare_adj_mat(adj_mat_trans)
|
|
|
|
adj_mat_train_trans = torch.sparse_coo_tensor(indices = edges_pos_trans[0].transpose(0, 1),
|
|
|
|
values=torch.ones(len(edges_pos_trans[0]), dtype=adj_mat_trans.dtype))
|
|
|
|
else:
|
|
|
|
adj_mat_train_trans = adj_mat_trans = \
|
|
|
|
edge_pos_trans = edge_neg_trans = None
|
|
|
|
|
|
|
|
return PreparedRelation(r.node_type_row, r.node_type_column,
|
|
|
|
adj_mat_train, adj_mat_trans_train,
|
|
|
|
edges_pos, edges_neg, edges_pos_trans, edges_neg_trans)
|
|
|
|
|
|
|
|
|
|
|
|
def prepare_training(data):
|
|
|
|
for (node_type_row, node_type_column), rels in data.relation_types:
|
|
|
|
for r in rels:
|
|
|
|
prep_relation_edges()
|