def combine_matrices_flat(full_relation, a_pos_heads, a_pos_tails, a_neg_heads,
a_neg_tails, ids, b_matrix, device):
'''
inputs:
a_heads: a dict of ID : head indices
a_tails: a dict of ID : tail indices
ids: IDs with which to access the indices of A
b_matrix: a matrix whose indices we want to include in output
returns:
out_matrix: matrix with indices & values of A as well as indices of B
valid_masks: a dict of id:indices that correspond to the indices for
each of the relations in A
'''
full_heads, full_tails = np.array([],
dtype=np.int32), np.array([],
dtype=np.int32)
for rel_id in ids:
full_heads = np.concatenate((full_heads, a_pos_heads[rel_id]))
full_heads = np.concatenate((full_heads, a_neg_heads[rel_id]))
full_tails = np.concatenate((full_tails, a_pos_tails[rel_id]))
full_tails = np.concatenate((full_tails, a_neg_tails[rel_id]))
indices = torch.LongTensor(np.vstack((full_heads, full_tails)))
values = torch.zeros((indices.shape[1], 1))
shape = (full_relation.entities[0].n_instances,
full_relation.entities[1].n_instances, 1)
full_a_matrix = SparseMatrix(indices=indices, values=values, shape=shape)
full_a_matrix = full_a_matrix.to(device).coalesce_()
b_idx_matrix = SparseMatrix.from_other_sparse_matrix(b_matrix, 1)
b_idx_matrix.values += 1
out_idx_matrix = b_idx_matrix + full_a_matrix
out_matrix = SparseMatrix.from_other_sparse_matrix(out_idx_matrix, 0)
for rel_id in ids:
rel_matrix = make_target_matrix(full_relation, a_pos_heads[rel_id],
a_pos_tails[rel_id],
a_neg_heads[rel_id],
a_neg_tails[rel_id], device)
rel_full_matrix = SparseMatrix.from_other_sparse_matrix(
out_idx_matrix, 1) + rel_matrix
out_matrix.values = torch.cat(
[out_matrix.values, rel_full_matrix.values], 1)
out_matrix.n_channels += 1
rel_idx_matrix = SparseMatrix.from_other_sparse_matrix(rel_matrix, 1)
rel_idx_matrix.values += 1
rel_idx_full_matrix = SparseMatrix.from_other_sparse_matrix(
out_idx_matrix, 1) + rel_idx_matrix
out_idx_matrix.values = torch.cat(
[out_idx_matrix.values, rel_idx_full_matrix.values], 1)
out_idx_matrix.n_channels += 1
masks = {}
for channel_i, rel_id in enumerate(ids):
masks[rel_id] = out_idx_matrix.values[:, channel_i +
1].nonzero().squeeze()
return out_matrix, masks
def make_flat_target_matrix(full_relation, rel_ids, pos_heads, pos_tails, neg_heads, neg_tails, device):
full_heads, full_tails = np.array([], dtype=np.int32), np.array([], dtype=np.int32)
for rel_id in rel_ids:
full_heads = np.concatenate((full_heads, pos_heads[rel_id]))
full_heads = np.concatenate((full_heads, neg_heads[rel_id]))
full_tails = np.concatenate((full_tails, pos_tails[rel_id]))
full_tails = np.concatenate((full_tails, neg_tails[rel_id]))
n_rels = len(rel_ids)
indices = torch.LongTensor(np.vstack((full_heads, full_tails)))
values = torch.zeros((indices.shape[1], n_rels))
shape = (full_relation.entities[0].n_instances,
full_relation.entities[1].n_instances, n_rels)
full_matrix = SparseMatrix(indices=indices, values=values, shape=shape)
full_matrix = full_matrix.to(device).coalesce_()
matrix_out = SparseMatrix.from_other_sparse_matrix(full_matrix, 0)
for rel_id in rel_ids:
rel_matrix = make_target_matrix(full_relation, pos_heads[rel_id],
pos_tails[rel_id], neg_heads[rel_id],
neg_tails[rel_id], device)
rel_matrix_full = SparseMatrix.from_other_sparse_matrix(full_matrix, 1) + rel_matrix
matrix_out.values = torch.cat([matrix_out.values, rel_matrix_full.values], 1)
matrix_out.n_channels += 1
return matrix_out
def test_broadcast_all(self):
zero_matrix = SparseMatrix.from_other_sparse_matrix(self.X,
n_channels=1)
out = zero_matrix.broadcast(torch.Tensor([5.]), "all")
'''
5ooo
5o5o
oooo
5oo5
'''
self.assertSameValues(out.values, np.array([[5, 5, 5, 5, 5]]).T)
def test_broadcast_row(self):
zero_matrix = SparseMatrix.from_other_sparse_matrix(self.X,
n_channels=1)
out = zero_matrix.broadcast(self.pooled, "row")
'''
1ooo
1o3o
oooo
1oo4
'''
self.assertSameValues(out.values, np.array([[1, 1, 3, 1, 4]]).T)
def test_broadcast_col(self):
zero_matrix = SparseMatrix.from_other_sparse_matrix(self.X,
n_channels=1)
out = zero_matrix.broadcast(self.pooled, "col")
'''
1ooo
2o2o
oooo
4oo4
'''
self.assertSameValues(out.values, np.array([[1, 2, 2, 4, 4]]).T)
def test_broadcast_diag(self):
zero_matrix = SparseMatrix.from_other_sparse_matrix(self.X,
n_channels=1)
out = zero_matrix.broadcast(torch.Tensor([5., 2]), "diag")
'''
5ooo
0o0o
oooo
0oo5
'''
self.assertSameValues(out.values,
np.array([[5, 0, 0, 0, 5], [2, 0, 0, 0, 2]]).T)
def forward(self, X_in, X_out, indices_identity, indices_trans):
'''
X_in: Source sparse tensor
X_out: Correpsonding sparse tensor for target relation
'''
self.logger.info("n_params: {}".format(self.n_params))
if type(X_out) == SparseMatrix:
Y = SparseMatrix.from_other_sparse_matrix(X_out, self.out_dim)
else:
Y = X_out.clone()
#TODO: can add a cache for input operations here
for i in range(self.n_params):
op_inp, op_out = self.all_ops[i]
weight = self.weights[i]
device = weight.device
if op_inp == None:
X_mul = torch.matmul(X_in, weight)
X_op_out = self.output_op(op_out, X_out, X_mul, device)
elif op_out == None:
X_op_inp = self.input_op(op_inp, X_in, device)
X_mul = torch.matmul(X_op_inp, weight)
X_op_out = X_mul
elif op_out[0] == "i":
# Identity
X_intersection_vals = X_in.gather_mask(indices_identity[0])
X_mul = X_intersection_vals @ weight
X_op_out = X_out.broadcast_from_mask(X_mul,
indices_identity[1],
device)
elif op_out[0] == "t":
# Transpose
X_T_intersection_vals = X_in.gather_transpose(indices_trans[0])
X_mul = X_T_intersection_vals @ weight
X_op_out = X_out.broadcast_from_mask(X_mul, indices_trans[1],
device)
else:
# Pool or Gather or Do Nothing
X_op_inp = self.input_op(op_inp, X_in, device)
# Multiply values by weight
X_mul = torch.matmul(X_op_inp, weight)
# Broadcast or Embed Diag or Transpose
X_op_out = self.output_op(op_out, X_out, X_mul, device)
#assert X_op_out.nnz() == X_out.nnz()
#assert Y.nnz() == X_out.nnz(), "Y: {}, X_out: {}".format(Y.nnz(), X_out.nnz())
#assert Y.nnz() == X_op_out.nnz(), "Y: {}, X_op_out: {}".format(Y.nnz(), X_op_out.nnz())
Y = Y + X_op_out
return Y
def select_features(data, schema, feats_type, target_ent):
'''
TODO: IMPLEMENT THIS
'''
# Select features for nodes
in_dims = {}
num_relations = len(schema.relations) - len(schema.entities)
if feats_type == 0:
# Keep all node attributes
pass
elif feats_type == 1:
# Set all non-target node attributes to zero
for ent_i in schema.entities:
if ent_i.id != target_ent:
# 10 dimensions for some reason
n_dim = 10
rel_id = num_relations + ent_i.id
data[rel_id] = SparseMatrix.from_other_sparse_matrix(
data[rel_id], n_dim)
'''
elif feats_type == 2:
# Set all non-target node attributes to one-hot vector
for i in range(0, len(features_list)):
if i != target_ent:
dim = features_list[i].shape[0]
indices = torch.arange(n_instances).unsqueeze(0).repeat(2, 1)
values = torch.FloatTensor(np.ones(dim))
features_list[i] = torch.sparse.FloatTensor(indices, values, torch.Size([dim, dim])).to(device)
elif feats_type == 3:
in_dims = [features.shape[0] for features in features_list]
for i in range(len(features_list)):
dim = features_list[i].shape[0]
indices = np.vstack((np.arange(dim), np.arange(dim)))
indices = torch.LongTensor(indices)
values = np.ones(dim)
features_list[i] = torch.sparse.FloatTensor(indices, values, torch.Size([dim, dim])).to(device)
'''
for rel_id in schema.relations:
in_dims[rel_id] = data[rel_id].n_channels
return data, in_dims
def load_data_flat(prefix,
use_node_attrs=True,
use_edge_data=True,
node_val='one'):
'''
Load data into one matrix with all relations, reproducing Maron 2019
The first [# relation types] channels are adjacency matrices,
while the next [sum of feature dimensions per entity type] channels have
node attributes on the relevant segment of their diagonals if use_node_attrs=True.
If node features aren't included, then ndoe_val is used instead.
'''
dl = data_loader(DATA_FILE_DIR + prefix)
total_n_nodes = dl.nodes['total']
entities = [Entity(0, total_n_nodes)]
relations = {0: Relation(0, [entities[0], entities[0]])}
schema = DataSchema(entities, relations)
# Sparse Matrix containing all data
data_full = sum(dl.links['data'].values()).tocoo()
data_diag = scipy.sparse.coo_matrix(
(np.ones(total_n_nodes),
(np.arange(total_n_nodes), np.arange(total_n_nodes))),
(total_n_nodes, total_n_nodes))
data_full += data_diag
data_full = SparseMatrix.from_scipy_sparse(data_full.tocoo()).zero_()
data_out = SparseMatrix.from_other_sparse_matrix(data_full, 0)
# Load up all edge data
for rel_id in sorted(dl.links['data'].keys()):
data_matrix = dl.links['data'][rel_id]
data_rel = SparseMatrix.from_scipy_sparse(data_matrix.tocoo())
if not use_edge_data:
# Use only adjacency information
data_rel.values = torch.ones(data_rel.values.shape)
data_rel_full = SparseMatrix.from_other_sparse_matrix(data_full,
1) + data_rel
data_out.values = torch.cat([data_out.values, data_rel_full.values], 1)
data_out.n_channels += 1
if use_node_attrs:
for ent_id, attr_matrix in dl.nodes['attr'].items():
start_i = dl.nodes['shift'][ent_id]
n_instances = dl.nodes['count'][ent_id]
if attr_matrix is None:
if node_val == 'zero':
attr_matrix = np.zeros((n_instances, 1))
elif node_val == 'rand':
attr_matrix = np.random.randn(n_instances, 1)
else:
attr_matrix = np.ones((n_instances, 1))
n_channels = attr_matrix.shape[1]
indices = torch.arange(start_i,
start_i + n_instances).unsqueeze(0).repeat(
2, 1)
data_rel = SparseMatrix(
indices=indices,
values=torch.FloatTensor(attr_matrix),
shape=np.array([total_n_nodes, total_n_nodes, n_channels]),
is_set=True)
data_rel_full = SparseMatrix.from_other_sparse_matrix(
data_full, n_channels) + data_rel
data_out.values = torch.cat(
[data_out.values, data_rel_full.values], 1)
data_out.n_channels += n_channels
data = SparseMatrixData(schema)
data[0] = data_out
return schema,\
data, \
dl
def load_data_flat(prefix,
use_node_attrs=True,
use_edge_data=True,
node_val='zero',
feats_type=0):
'''
Load data into one matrix with all relations, reproducing Maron 2019
The first [# relation types] channels are adjacency matrices,
while the next [sum of feature dimensions per entity type] channels have
node attributes on the relevant segment of their diagonals if use_node_attrs=True.
If node features aren't included, then ndoe_val is used instead.
'''
dl = data_loader(DATA_FILE_DIR + prefix)
total_n_nodes = dl.nodes['total']
entities = [Entity(0, total_n_nodes)]
relations = {0: Relation(0, [entities[0], entities[0]])}
schema = DataSchema(entities, relations)
# Sparse Matrix containing all data
data_full = sum(dl.links['data'].values()).tocoo()
data_diag = scipy.sparse.coo_matrix(
(np.ones(total_n_nodes),
(np.arange(total_n_nodes), np.arange(total_n_nodes))),
(total_n_nodes, total_n_nodes))
data_full += data_diag
data_full = SparseMatrix.from_scipy_sparse(data_full.tocoo()).zero_()
data_out = SparseMatrix.from_other_sparse_matrix(data_full, 0)
# Load up all edge data
for rel_id in sorted(dl.links['data'].keys()):
data_matrix = dl.links['data'][rel_id]
data_rel = SparseMatrix.from_scipy_sparse(data_matrix.tocoo())
if not use_edge_data:
# Use only adjacency information
data_rel.values = torch.ones(data_rel.values.shape)
data_rel_full = SparseMatrix.from_other_sparse_matrix(data_full,
1) + data_rel
data_out.values = torch.cat([data_out.values, data_rel_full.values], 1)
data_out.n_channels += 1
target_entity = 0
if use_node_attrs:
for ent_id, attr_matrix in dl.nodes['attr'].items():
start_i = dl.nodes['shift'][ent_id]
n_instances = dl.nodes['count'][ent_id]
if attr_matrix is None:
if node_val == 'zero':
attr_matrix = np.zeros((n_instances, 1))
elif node_val == 'rand':
attr_matrix = np.random.randn(n_instances, 1)
else:
attr_matrix = np.ones((n_instances, 1))
if feats_type == 1 and ent_id != target_entity:
# To keep same behaviour as non-LGNN model, use 10 dimensions
attr_matrix = np.zeros((n_instances, 10))
n_channels = attr_matrix.shape[1]
indices = torch.arange(start_i,
start_i + n_instances).unsqueeze(0).repeat(
2, 1)
data_rel = SparseMatrix(
indices=indices,
values=torch.FloatTensor(attr_matrix),
shape=np.array([total_n_nodes, total_n_nodes, n_channels]),
is_set=True)
data_rel_full = SparseMatrix.from_other_sparse_matrix(
data_full, n_channels) + data_rel
data_out.values = torch.cat(
[data_out.values, data_rel_full.values], 1)
data_out.n_channels += n_channels
data = SparseMatrixData(schema)
data[0] = data_out
n_outputs = total_n_nodes
n_output_classes = dl.labels_train['num_classes']
schema_out = DataSchema([entities[target_entity]], [
Relation(0, [entities[target_entity], entities[target_entity]],
is_set=True)
])
data_target = SparseMatrixData(schema_out)
data_target[0] = SparseMatrix(
indices=torch.arange(n_outputs, dtype=torch.int64).repeat(2, 1),
values=torch.zeros([n_outputs, n_output_classes]),
shape=(n_outputs, n_outputs, n_output_classes),
is_set=True)
labels = np.zeros((dl.nodes['count'][0], dl.labels_train['num_classes']),
dtype=int)
val_ratio = 0.2
train_idx = np.nonzero(dl.labels_train['mask'])[0]
np.random.shuffle(train_idx)
split = int(train_idx.shape[0] * val_ratio)
val_idx = train_idx[:split]
train_idx = train_idx[split:]
train_idx = np.sort(train_idx)
val_idx = np.sort(val_idx)
test_idx = np.nonzero(dl.labels_test['mask'])[0]
labels[train_idx] = dl.labels_train['data'][train_idx]
labels[val_idx] = dl.labels_train['data'][val_idx]
if prefix != 'IMDB':
labels = labels.argmax(axis=1)
train_val_test_idx = {}
train_val_test_idx['train_idx'] = train_idx
train_val_test_idx['val_idx'] = val_idx
train_val_test_idx['test_idx'] = test_idx
return schema,\
schema_out, \
data, \
data_target, \
labels,\
train_val_test_idx,\
dl