def backward(ctx, grad_indexC, grad_valueC):
m, k = ctx.m, ctx.k
n = ctx.n
indexA, valueA, indexB, valueB, indexC = ctx.saved_tensors
grad_valueA = grad_valueB = None
if not grad_valueC.is_cuda:
if ctx.needs_input_grad[1] or ctx.needs_input_grad[1]:
grad_valueC = grad_valueC.clone()
if ctx.needs_input_grad[1]:
grad_valueA = torch_sparse.spspmm_cpu.spspmm_bw(
indexA, indexC.detach(), grad_valueC, indexB.detach(),
valueB, m, k)
if ctx.needs_input_grad[3]:
indexA, valueA = transpose(indexA, valueA, m, k)
indexC, grad_valueC = transpose(indexC, grad_valueC, m, n)
grad_valueB = torch_sparse.spspmm_cpu.spspmm_bw(
indexB, indexA.detach(), valueA, indexC.detach(),
grad_valueC, k, n)
else:
if ctx.needs_input_grad[1]:
grad_valueA = torch_sparse.spspmm_cuda.spspmm_bw(
indexA, indexC.detach(), grad_valueC.clone(),
indexB.detach(), valueB, m, k)
if ctx.needs_input_grad[3]:
indexA_T, valueA_T = transpose(indexA, valueA, m, k)
grad_indexB, grad_valueB = mm(indexA_T, valueA_T, indexC,
grad_valueC, k, m, n)
grad_valueB = lift(grad_indexB, grad_valueB, indexB, n)
return None, grad_valueA, None, grad_valueB, None, None, None
def is_undirected(edge_index, edge_attr=None, num_nodes=None):
r"""Returns :obj:`True` if the graph given by :attr:`edge_index` is
undirected.
Args:
edge_index (LongTensor): The edge indices.
edge_attr (Tensor, optional): Edge weights or multi-dimensional
edge features. (default: :obj:`None`)
num_nodes (int, optional): The number of nodes, *i.e.*
:obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
:rtype: bool
"""
num_nodes = maybe_num_nodes(edge_index, num_nodes)
edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes,
num_nodes)
if edge_attr is None:
undirected_edge_index = to_undirected(edge_index, num_nodes=num_nodes)
return edge_index.size(1) == undirected_edge_index.size(1)
else:
edge_index_t, edge_attr_t = transpose(edge_index,
edge_attr,
num_nodes,
num_nodes,
coalesced=True)
index_symmetric = torch.all(edge_index == edge_index_t)
attr_symmetric = torch.all(edge_attr == edge_attr_t)
return index_symmetric and attr_symmetric
def loop_sparse_attention_centrality(self, attention, idx):
# O(N) implementation
batch, groups, npoints, neighbors = attention.size()
idx_tag = torch.tensor([[i] * neighbors for i in range(npoints)],
device='cuda').flatten().unsqueeze(0)
mtrx = torch.tensor([[1.] * npoints], device='cuda').T
score = []
for i in range(batch):
idx_flatten = idx[i].flatten().unsqueeze(0) # NK
index = torch.cat([idx_tag, idx_flatten], dim=0)
score_group = []
for j in range(groups):
attention_flatten = attention[i][j].flatten()
index_s, value_s = coalesce(index, attention_flatten, npoints,
npoints)
index_t, value_t = transpose(index_s, value_s, npoints,
npoints)
out = spmm(index_t, value_t, npoints, npoints, mtrx)
score_group.append(out)
if j == groups - 1:
score.append(torch.cat(score_group, dim=1).unsqueeze(0))
if i == batch - 1:
final_score = torch.cat(score, dim=0)
# fullnl instance
final_score = final_score.unsqueeze(3).permute(0, 2, 3, 1)
idx_value, idx_score = final_score.topk(
k=neighbors, dim=3) # B, G, 1, N -> B, G, 1, K'
return idx_value, idx_score
def prepare_groups(self, features, mask):
#tensor_mask = torch.from_numpy(mask)
#self.groups = torch.clamp(self.groups[tensor_mask, :], 0, 1).transpose_(1, 0)
#padding_a = features.shape[1] - self.groups.shape[0]
#if padding_a > 0:
# padding_a = ConstantPad2d((0, 0, 0, padding_a), 0)
# self.groups = padding_a(self.groups)
#if not self.sparse_groups.is_coalesced():
# self.sparse_groups = self.sparse_groups.coalesce()
where_mask = torch.from_numpy(np.argwhere(mask).squeeze()).to(
self.sparse_groups.device)
value_mask = self.isin(self.sparse_groups._indices()[0, :], where_mask)
values = self.sparse_groups._values()[value_mask]
indicies = self.sparse_groups._indices()[:, value_mask]
values = torch.clamp(values, 0, 1)
new_cols = indicies[1, :]
# Rescale row_idx to slice
temp = indicies[0, :]
d = {j: i for i, j in enumerate(temp.unique().cpu().numpy())}
new_rows = temp.cpu().apply_(lambda x: d[x]).to(
self.sparse_groups.device)
indicies = torch.stack([new_rows, new_cols])
from torch_sparse import transpose
indicies, values = transpose(indicies, values,
self.sparse_groups.shape[0],
self.sparse_groups.shape[1])
self.sparse_groups = torch.sparse.FloatTensor(
indicies, values, (features.shape[1], mask.sum()))
def test_transpose():
row = torch.tensor([1, 0, 1, 0, 2, 1])
col = torch.tensor([0, 1, 1, 1, 0, 0])
index = torch.stack([row, col], dim=0)
value = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
index, value = transpose(index, value, m=3, n=2)
assert index.tolist() == [[0, 0, 1, 1], [1, 2, 0, 1]]
assert value.tolist() == [[7, 9], [5, 6], [6, 8], [3, 4]]
def test_transpose_matrix(dtype, device):
row = torch.tensor([1, 0, 1, 2], device=device)
col = torch.tensor([0, 1, 1, 0], device=device)
index = torch.stack([row, col], dim=0)
value = tensor([1, 2, 3, 4], dtype, device)
index, value = transpose(index, value, m=3, n=2)
assert index.tolist() == [[0, 0, 1, 1], [1, 2, 0, 1]]
assert value.tolist() == [1, 4, 2, 3]
def transpose(self, transpose_mask=None):
indices_out, values_out = torch_sparse.transpose(self.indices, self.values,
self.n, self.m, coalesced=True)
if transpose_mask is not None:
values_out = (transpose_mask.float() * values_out.T).T
shape_out = (self.m, self.n, self.n_channels)
return SparseMatrix(indices_out, values_out, shape_out, self.indices_diag, self.is_set)
def backward(ctx, grad_indexC, grad_valueC):
m, k, n = ctx.m, ctx.k, ctx.n
indexA, valueA, indexB, valueB, indexC = ctx.saved_tensors
grad_valueA = grad_valueB = None
if ctx.needs_input_grad[1]:
indexB_T, valueB_T = transpose(indexB, valueB, k, n)
grad_indexA, grad_valueA = mm(indexC, grad_valueC, indexB_T,
valueB_T, m, n, k)
grad_valueA = lift(grad_indexA, grad_valueA, indexA, k)
if ctx.needs_input_grad[3]:
indexA_T, valueA_T = transpose(indexA, valueA, m, k)
grad_indexB, grad_valueB = mm(indexA_T, valueA_T, indexC,
grad_valueC, k, m, n)
grad_valueB = lift(grad_indexB, grad_valueB, indexB, n)
return None, grad_valueA, None, grad_valueB, None, None, None
def batched_transpose(adj, value, m=None, n=None):
"""
Args:
adj: Tensor or list of Tensor
value: Tensor [num_edges, ]
m: int
n: int
"""
if isinstance(adj, Tensor):
m = maybe_num_nodes(adj[0], m)
n = maybe_num_nodes(adj[1], n)
return transpose(adj, value, m, n)
else: # adj is a list of Tensor
adj_ = [None] * value.shape[1]
vs = torch.zeros(value.shape)
m = max([maybe_num_nodes(a_[0], m) for a_ in adj])
n = max([maybe_num_nodes(a_[1], n) for a_ in adj])
for j in range(len(adj)):
adj_[j], vs[:, j] = transpose(adj[j], value[:, j], m, n)
return adj_, vs
def test_sparse(self):
index = torch.tensor([[1, 0, 1, 0, 2, 1],
[0, 1, 1, 1, 0, 0]])
value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]])
index, value = torch_sparse.transpose(index, value, 3, 2)
test_index = torch.LongTensor([[0, 0, 1, 1],
[1, 2, 0, 1]])
test_value = torch.Tensor([[7., 9.], [5., 6.], [6., 8.], [3., 4.]])
self.assertTrue(torch.all(torch.eq(test_index, index)))
self.assertTrue(torch.all(torch.eq(test_value, value)))
def StAS(index_A, value_A, index_S, value_S, device, N, kN):
r"""StAS: a function which returns new edge weights for the pooled graph using the formula S^{T}AS"""
index_A, value_A = coalesce(index_A, value_A, m=N, n=N)
index_S, value_S = coalesce(index_S, value_S, m=N, n=kN)
index_B, value_B = spspmm(index_A, value_A, index_S, value_S, N, N, kN)
index_St, value_St = transpose(index_S, value_S, N, kN)
index_B, value_B = coalesce(index_B, value_B, m=N, n=kN)
index_E, value_E = spspmm(index_St.cpu(), value_St.cpu(), index_B.cpu(), value_B.cpu(), kN, N, kN)
return index_E.to(device), value_E.to(device)
def add_reverse_edge_index(self, edge_index_dict) -> None:
reverse_edge_index_dict = {}
for metapath in edge_index_dict:
if is_negative(metapath) or edge_index_dict[metapath] == None:
continue
reverse_metapath = self.reverse_metapath_name(
metapath, edge_index_dict)
reverse_edge_index_dict[reverse_metapath] = transpose(
index=edge_index_dict[metapath],
value=None,
m=self.num_nodes_dict[metapath[0]],
n=self.num_nodes_dict[metapath[-1]])[0]
edge_index_dict.update(reverse_edge_index_dict)
def rebuild_features_average(self, features, mask, target_edges):
self.prepare_groups(features, mask)
if not self.sparse_groups.is_coalesced():
self.sparse_groups = self.sparse_groups.coalesce()
idxs, values = transpose(self.sparse_groups._indices(),
self.sparse_groups._values(),
self.sparse_groups.shape[0],
self.sparse_groups.shape[1])
fe = spmm(idxs, values, self.sparse_groups.shape[1],
self.sparse_groups.shape[0],
features.squeeze(-1).T).T
# fe = torch.matmul(features.squeeze(-1), self.groups)
#occurrences = torch.sum(self.groups, 0).expand(fe.shape)
occurrences = torch.sparse.sum(self.sparse_groups,
0).to_dense().expand(fe.shape)
fe = fe / occurrences
padding_b = target_edges - fe.shape[1]
if padding_b > 0:
padding_b = ConstantPad2d((0, padding_b, 0, 0), 0)
fe = padding_b(fe)
return fe
else:
pyg_dataset = PygNodePropPredDataset(name='ogbn-mag',
root='../dataset')
pyg_data = pyg_dataset[0]
# ('author', 'affiliated_with', 'institution'),
# ('author', 'writes', 'paper'),
# ('paper', 'cites', 'paper'),
# ('paper', 'has_topic', 'field_of_study')
# reverse edges (relations) in the heterogeneous graph.
pyg_data.edge_index_dict[(
'institution', 'rev_affiliated_with',
'author')] = transpose(pyg_data.edge_index_dict[('author',
'affiliated_with',
'institution')],
None,
m=pyg_data.num_nodes_dict['author'],
n=pyg_data.num_nodes_dict['institution'])[0]
pyg_data.edge_index_dict[('paper', 'rev_writes',
'author')] = transpose(
pyg_data.edge_index_dict[('author',
'writes',
'paper')],
None,
m=pyg_data.num_nodes_dict['author'],
n=pyg_data.num_nodes_dict['paper'])[0]
pyg_data.edge_index_dict[('paper', 'rev_cites', 'paper')] = transpose(
pyg_data.edge_index_dict[('paper', 'cites', 'paper')],
None,
def parallel_sparse_attention_centrality(self, attention, idx):
# O(N) implementation
cnt = 0
#MEMCHK#
tic = time.time()
cnt, tic = memchk(cnt, string="AC Start", level=8, tic=tic)
#MEMCHK#
batch, groups, npoints, neighbors = attention.size()
idx_arange = torch.tensor([[npoints * i] * npoints * neighbors
for i in range(batch * groups)],
device='cuda').flatten() # NK -> 11NK
idx_self = torch.tensor(
[[i] * neighbors for i in range(npoints)],
device='cuda').flatten().unsqueeze(0).unsqueeze(0) # NK -> 11NK
idx_self = idx_self.repeat(batch, groups, 1) # BGNK
mtrx = torch.tensor([[1.] * npoints * batch * groups], device='cuda').T
idx_neighbor_flatten = idx.repeat(1, groups, 1, 1).flatten() # BGNK
idx_neighbor_flatten = (idx_arange + idx_neighbor_flatten).unsqueeze(
0) # 1, BGNK
idx_self_flatten = idx_self.flatten() # BGNK
idx_self_flatten = (idx_arange + idx_self_flatten).unsqueeze(
0) # 1, BGNK
index = torch.cat([idx_self_flatten, idx_neighbor_flatten],
dim=0) # 2, BGNK
attention_flatten = attention.flatten() # BGNK
#MEMCHK#
cnt, tic = memchk(cnt, string="CoalesceStart", level=8, tic=tic)
#MEMCHK#
index_s, value_s = coalesce(index, attention_flatten,
npoints * batch * groups,
npoints * batch * groups)
index_t, value_t = transpose(index_s, value_s,
npoints * batch * groups,
npoints * batch * groups)
out = spmm(index_t, value_t, npoints * batch * groups,
npoints * batch * groups, mtrx)
#MEMCHK#
cnt, tic = memchk(cnt, string="CoalesceEnd", level=8, tic=tic)
#MEMCHK#
out = out.view(batch, groups, npoints, 1)
score = out.permute(0, 1, 3, 2)
idx_value, idx_score = score.topk(k=neighbors,
dim=3) # B, G, 1, N -> B, G, 1, K'
#MEMCHK#
cnt, tic = memchk(cnt, string="AC End", level=8, tic=tic)
#MEMCHK#
return idx_value, idx_score
def mgpool(x, pos, edge_index, batch, mask=None):
adj_values = torch.ones(edge_index.shape[1]).cuda()
cluster = graclus(edge_index)
cluster, perm = consecutive_cluster(cluster)
index = torch.stack([cluster, torch.arange(0, x.shape[0]).cuda()], dim=0)
values = torch.ones(cluster.shape[0], dtype=torch.float).cuda()
uniq, inv, counts = torch.unique(cluster,
return_inverse=True,
return_counts=True)
newsize = uniq.shape[0]
origsize = x.shape[0]
new_batch = pool_batch(perm, batch)
# Compute random walk graph laplacian:
laplacian_index, laplacian_weights = get_laplacian(edge_index,
normalization='rw')
laplacian_index, laplacian_weights = torch_sparse.coalesce(
laplacian_index, laplacian_weights, m=origsize, n=origsize)
index, values = torch_sparse.coalesce(index, values, m=newsize,
n=origsize) # P^T matrix
new_feat = torch_sparse.spmm(index,
values,
m=newsize,
n=origsize,
matrix=x) # P^T X
new_pos = torch_sparse.spmm(index,
values,
m=newsize,
n=origsize,
matrix=pos) # P^T POS
new_adj, new_adj_val = torch_sparse.spspmm(index,
values,
edge_index,
adj_values,
m=newsize,
k=origsize,
n=origsize,
coalesced=True) # P^T A
index, values = torch_sparse.transpose(index,
values,
m=newsize,
n=origsize,
coalesced=True) # P
new_adj, new_adj_val = torch_sparse.spspmm(new_adj,
new_adj_val,
index,
values,
m=newsize,
k=origsize,
n=newsize,
coalesced=True) # (P^T A) P
# Precompute QP :
values = torch.ones(cluster.shape[0], dtype=torch.float).cuda()
index, values = torch_sparse.spspmm(laplacian_index,
laplacian_weights,
index,
values,
m=origsize,
k=origsize,
n=newsize,
coalesced=True)
return new_adj, new_feat, new_pos, new_batch, index, values, origsize, newsize
def process(self):
import pandas as pd
data = HeteroData()
# Get author labels.
path = osp.join(self.raw_dir, 'id_author.txt')
author = pd.read_csv(path,
sep='\t',
names=['idx', 'name'],
index_col=1)
path = osp.join(self.raw_dir, 'label',
'googlescholar.8area.author.label.txt')
df = pd.read_csv(path, sep=' ', names=['name', 'y'])
df = df.join(author, on='name')
data['author'].y = torch.from_numpy(df['y'].values) - 1
data['author'].y_index = torch.from_numpy(df['idx'].values)
# Get venue labels.
path = osp.join(self.raw_dir, 'id_conf.txt')
venue = pd.read_csv(path, sep='\t', names=['idx', 'name'], index_col=1)
path = osp.join(self.raw_dir, 'label',
'googlescholar.8area.venue.label.txt')
df = pd.read_csv(path, sep=' ', names=['name', 'y'])
df = df.join(venue, on='name')
data['venue'].y = torch.from_numpy(df['y'].values) - 1
data['venue'].y_index = torch.from_numpy(df['idx'].values)
# Get paper<->author connectivity.
path = osp.join(self.raw_dir, 'paper_author.txt')
paper_author = pd.read_csv(path, sep='\t', header=None)
paper_author = torch.from_numpy(paper_author.values)
paper_author = paper_author.t().contiguous()
M, N = int(paper_author[0].max() + 1), int(paper_author[1].max() + 1)
paper_author, _ = coalesce(paper_author, None, M, N)
author_paper, _ = transpose(paper_author, None, M, N)
data['paper'].num_nodes = M
data['author'].num_nodes = N
data['paper', 'written_by', 'author'].edge_index = paper_author
data['author', 'writes', 'paper'].edge_index = author_paper
# Get paper<->venue connectivity.
path = osp.join(self.raw_dir, 'paper_conf.txt')
paper_venue = pd.read_csv(path, sep='\t', header=None)
paper_venue = torch.from_numpy(paper_venue.values)
paper_venue = paper_venue.t().contiguous()
M, N = int(paper_venue[0].max() + 1), int(paper_venue[1].max() + 1)
paper_venue, _ = coalesce(paper_venue, None, M, N)
venue_paper, _ = transpose(paper_venue, None, M, N)
data['venue'].num_nodes = N
data['paper', 'published_in', 'venue'].edge_index = paper_venue
data['venue', 'publishes', 'paper'].edge_index = venue_paper
if self.pre_transform is not None:
data = self.pre_transform(data)
torch.save(self.collate([data]), self.processed_paths[0])
def process(self):
# Get author labels.
path = osp.join(self.raw_dir, 'id_author.txt')
author = pandas.read_csv(path,
sep='\t',
names=['idx', 'name'],
index_col=1)
path = osp.join(self.raw_dir, 'label',
'googlescholar.8area.author.label.txt')
df = pandas.read_csv(path, sep=' ', names=['name', 'y'])
df = df.join(author, on='name')
author_y = torch.from_numpy(df['y'].values) - 1
author_y_index = torch.from_numpy(df['idx'].values)
# Get venue labels.
path = osp.join(self.raw_dir, 'id_conf.txt')
venue = pandas.read_csv(path,
sep='\t',
names=['idx', 'name'],
index_col=1)
path = osp.join(self.raw_dir, 'label',
'googlescholar.8area.venue.label.txt')
df = pandas.read_csv(path, sep=' ', names=['name', 'y'])
df = df.join(venue, on='name')
venue_y = torch.from_numpy(df['y'].values) - 1
venue_y_index = torch.from_numpy(df['idx'].values)
# Get paper<->author connectivity.
path = osp.join(self.raw_dir, 'paper_author.txt')
paper_author = pandas.read_csv(path, sep='\t', header=None)
paper_author = torch.from_numpy(paper_author.values)
paper_author = paper_author.t().contiguous()
M, N = int(paper_author[0].max() + 1), int(paper_author[1].max() + 1)
paper_author, _ = coalesce(paper_author, None, M, N)
author_paper, _ = transpose(paper_author, None, M, N)
# Get paper<->venue connectivity.
path = osp.join(self.raw_dir, 'paper_conf.txt')
paper_venue = pandas.read_csv(path, sep='\t', header=None)
paper_venue = torch.from_numpy(paper_venue.values)
paper_venue = paper_venue.t().contiguous()
M, N = int(paper_venue[0].max() + 1), int(paper_venue[1].max() + 1)
paper_venue, _ = coalesce(paper_venue, None, M, N)
venue_paper, _ = transpose(paper_venue, None, M, N)
data = Data(
edge_index_dict={
('paper', 'written by', 'author'): paper_author,
('author', 'wrote', 'paper'): author_paper,
('paper', 'published in', 'venue'): paper_venue,
('venue', 'published', 'paper'): venue_paper,
},
y_dict={
'author': author_y,
'venue': venue_y,
},
y_index_dict={
'author': author_y_index,
'venue': venue_y_index,
},
num_nodes_dict={
'paper': int(paper_author[0].max()) + 1,
'author': author.shape[0],
'venue': venue.shape[0],
},
)
if self.pre_transform is not None:
data = self.pre_transform(data)
torch.save(self.collate([data]), self.processed_paths[0])
def gather_transpose(self, mask):
indices_out, values_out = torch_sparse.transpose(self.indices, self.values,
self.n, self.m, coalesced=True)
return values_out[mask]
def main():
parser = argparse.ArgumentParser(description='OGBN-MAG (MetaPath2Vec)')
parser.add_argument('--device', type=int, default=0)
parser.add_argument('--embedding_dim', type=int, default=128)
parser.add_argument('--walk_length', type=int, default=64)
parser.add_argument('--context_size', type=int, default=7)
parser.add_argument('--walks_per_node', type=int, default=5)
parser.add_argument('--num_negative_samples', type=int, default=5)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--lr', type=float, default=0.01)
parser.add_argument('--epochs', type=int, default=5)
parser.add_argument('--log_steps', type=int, default=100)
args = parser.parse_args()
device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
dataset = PygNodePropPredDataset('ogbn-mag')
data = dataset[0]
# We need to add reverse edges to the heterogeneous graph.
data.edge_index_dict[('institution', 'employs', 'author')] = transpose(
data.edge_index_dict[('author', 'affiliated_with', 'institution')],
None,
m=data.num_nodes_dict['author'],
n=data.num_nodes_dict['institution'])[0]
data.edge_index_dict[('paper', 'written_by', 'author')] = transpose(
data.edge_index_dict[('author', 'writes', 'paper')],
None,
m=data.num_nodes_dict['author'],
n=data.num_nodes_dict['paper'])[0]
data.edge_index_dict[('field_of_study', 'contains', 'paper')] = transpose(
data.edge_index_dict[('paper', 'has_topic', 'field_of_study')],
None,
m=data.num_nodes_dict['paper'],
n=data.num_nodes_dict['field_of_study'])[0]
print(data)
metapath = [
('author', 'writes', 'paper'),
('paper', 'has_topic', 'field_of_study'),
('field_of_study', 'contains', 'paper'),
('paper', 'written_by', 'author'),
('author', 'affiliated_with', 'institution'),
('institution', 'employs', 'author'),
('author', 'writes', 'paper'),
('paper', 'cites', 'paper'),
('paper', 'written_by', 'author'),
]
model = MetaPath2Vec(data.edge_index_dict,
embedding_dim=128,
metapath=metapath,
walk_length=64,
context_size=7,
walks_per_node=5,
num_negative_samples=5,
sparse=True).to(device)
loader = model.loader(batch_size=128, shuffle=True, num_workers=4)
optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)
model.train()
for epoch in range(1, args.epochs + 1):
for i, (pos_rw, neg_rw) in enumerate(loader):
optimizer.zero_grad()
loss = model.loss(pos_rw.to(device), neg_rw.to(device))
loss.backward()
optimizer.step()
if (i + 1) % args.log_steps == 0:
print(f'Epoch: {epoch:02d}, Step: {i+1:03d}/{len(loader)}, '
f'Loss: {loss:.4f}')
if (i + 1) % 1000 == 0: # Save model every 1000 steps.
save_embedding(model)
save_embedding(model)
def transpose(self):
(tidxs, tvals) = transpose(self.idxs, self.vals, self.shape[0], self.shape[1])
return SpTensor(tidxs, tvals, (self.shape[1], self.shape[0]))
def degrees(M):
signal = torch.reshape(torch.ones(M.shape[0]), (-1, 1))
index, values = torch_sparse.transpose(M.index, M.values, M.shape[0],
M.shape[1])
return torch.ravel(
torch_sparse.spmm(index, values, M.shape[1], M.shape[0], signal))
# 这里是按照行优先的顺序 # 按行排序索引并删除重复项。重复的条目使用 torch_scatter 的聚合函数计算 # 即存在空行或者空列的情况，合并稀疏矩阵重复位置处的数值 # 合并后只有4个位置处有数值 index = torch.tensor([[1, 0, 1, 0, 2, 1], [0, 1, 1, 1, 0, 0]]) value = torch.Tensor([[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]]) index, value = coalesce(index, value, m=3, n=2) print(index, value, sep='\n') # 聚合函数使用 平均值 index, value = coalesce(index, value, m=3, n=2, op='mean') print(index, value, sep='\n') # 转置，聚合函数默认为 add index, value = transpose(index, value, 3, 2) print(index, value, sep='\n') # 稀疏矩阵 * 稠密矩阵 # index 表示稀疏矩阵的坐标，value表示对应的值 # 第一行是行坐标，第二行是列坐标 index = torch.tensor([[0, 0, 1, 2, 2], [0, 2, 1, 0, 1]]) value = torch.Tensor([1, 2, 4, 1, 3]) matrix = torch.Tensor([[1, 4], [2, 5], [3, 6]]) out = spmm(index, value, 3, matrix) print(out) # 稀疏矩阵 * 稀疏矩阵 indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]]) valueA = torch.Tensor([1, 2, 3, 4, 5])
