def process(self, data):
recursive = '_recursive' if self.recursive else ''
filename = f'part_data_{self.num_parts}{recursive}.pt'
path = osp.join(self.save_dir or '', filename)
if self.save_dir is not None and osp.exists(path):
data, partptr, perm = torch.load(path)
else:
data = copy.copy(data)
num_nodes = data.num_nodes
(row, col), edge_attr = data.edge_index, data.edge_attr
adj = SparseTensor(row=row, col=col, value=edge_attr)
adj, partptr, perm = adj.partition(self.num_parts, self.recursive)
for key, item in data:
if item.size(0) == num_nodes:
data[key] = item[perm]
data.edge_index = None
data.edge_attr = None
data.adj = adj
if self.save_dir is not None:
torch.save((data, partptr, perm), path)
self.data = data
self.perm = perm
self.partptr = partptr
def __init__(self,
data,
num_parts,
recursive=False,
save_dir=None,
log=True):
assert data.edge_index is not None
recursive_str = '_recursive' if recursive else ''
filename = f'partition_{num_parts}{recursive_str}.pt'
path = osp.join(save_dir or '', filename)
if save_dir is not None and osp.exists(path):
adj, partptr, perm = torch.load(path)
else:
if log: # pragma: no cover
print('Compute METIS partitioning...')
(row, col), edge_attr = data.edge_index, data.edge_attr
adj = SparseTensor(row=row, col=col, value=edge_attr)
adj, partptr, perm = adj.partition(num_parts, recursive)
if save_dir is not None:
torch.save((adj, partptr, perm), path)
if log: # pragma: no cover
print('Done!')
self.data = self.__permute_data__(data, perm, adj)
self.partptr = partptr
self.perm = perm
def __init__(self, data, num_parts: int, recursive: bool = False,
save_dir: Optional[str] = None, log: bool = True):
assert data.edge_index is not None
self.num_parts = num_parts
recursive_str = '_recursive' if recursive else ''
filename = f'partition_{num_parts}{recursive_str}.pt'
path = osp.join(save_dir or '', filename)
if save_dir is not None and osp.exists(path):
adj, partptr, perm = torch.load(path)
else:
if log: # pragma: no cover
print('Computing METIS partitioning...')
N, E = data.num_nodes, data.num_edges
adj = SparseTensor(
row=data.edge_index[0], col=data.edge_index[1],
value=torch.arange(E, device=data.edge_index.device),
sparse_sizes=(N, N))
adj, partptr, perm = adj.partition(num_parts, recursive)
if save_dir is not None:
torch.save((adj, partptr, perm), path)
if log: # pragma: no cover
print('Done!')
self.data = self.__permute_data__(data, perm, adj)
self.partptr = partptr
self.perm = perm
def process(self):
recursive = '_recursive' if self.recursive else ''
filename = f'part_data_{self.num_parts}{recursive}.pt'
path = osp.join(self.dataset.processed_dir, filename)
if self.save and osp.exists(path):
data, partptr, perm = torch.load(path)
else:
data = copy.copy(self.dataset.get(0))
(row, col), edge_attr = data.edge_index, data.edge_attr
adj = SparseTensor(row=row, col=col, value=edge_attr,
is_sorted=True)
adj, partptr, perm = adj.partition(self.num_parts, self.recursive)
for key, item in data:
if item.size(0) == data.num_nodes:
data[key] = item[perm]
data.edge_index = None
data.edge_attr = None
data.adj = adj
if self.save:
torch.save((data, partptr, perm), path)
self.__data__ = data
self.__perm__ = perm
self.__partptr__ = partptr
def __init__(self, data, num_parts, recursive=False, save_dir=None):
assert data.edge_index is not None
recursive_str = '_recursive' if recursive else ''
filename = f'partition_{num_parts}{recursive_str}.pt'
path = osp.join(save_dir or '', filename)
if save_dir is not None and osp.exists(path):
adj, partptr, perm = torch.load(path)
else:
(row, col), edge_attr = data.edge_index, data.edge_attr
adj = SparseTensor(row=row, col=col, value=edge_attr)
adj, partptr, perm = adj.partition(num_parts, recursive)
if save_dir is not None:
torch.save((adj, partptr, perm), path)
self.data = self.permute_data(data, perm, adj)
self.partptr = partptr
self.perm = perm
