Spaces:
Running on Zero
Running on Zero
File size: 2,846 Bytes
6fe4366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import torch
import torch_geometric.utils
from torch_geometric.utils import to_dense_adj, to_dense_batch
def to_dense(x, edge_index, edge_attr, batch):
X, node_mask = to_dense_batch(x=x, batch=batch)
edge_index, edge_attr = torch_geometric.utils.remove_self_loops(
edge_index, edge_attr
)
max_num_nodes = X.size(1)
E = to_dense_adj(
edge_index=edge_index,
batch=batch,
edge_attr=edge_attr,
max_num_nodes=max_num_nodes,
)
E = encode_no_edge(E)
return PlaceHolder(X=X, E=E, y=None), node_mask
def encode_no_edge(E):
assert len(E.shape) == 4
if E.shape[-1] == 0:
return E
no_edge = torch.sum(E, dim=3) == 0
first_elt = E[:, :, :, 0]
first_elt[no_edge] = 1
E[:, :, :, 0] = first_elt
diag = (
torch.eye(E.shape[1], dtype=torch.bool).unsqueeze(0).expand(E.shape[0], -1, -1)
)
E[diag] = 0
return E
class PlaceHolder:
def __init__(self, X, E, y):
self.X = X
self.E = E
self.y = y
def type_as(self, x: torch.Tensor):
"""Changes the device and dtype of X, E, y."""
self.X = self.X.type_as(x)
self.E = self.E.type_as(x)
self.y = self.y.type_as(x)
return self
def to_device(self, device):
"""Changes the device and dtype of X, E, y."""
self.X = self.X.to(device)
self.E = self.E.to(device)
self.y = self.y.to(device) if self.y is not None else None
return self
def mask(self, node_mask, collapse=False):
x_mask = node_mask.unsqueeze(-1) # bs, n, 1
e_mask1 = x_mask.unsqueeze(2) # bs, n, 1, 1
e_mask2 = x_mask.unsqueeze(1) # bs, 1, n, 1
if collapse:
self.X = torch.argmax(self.X, dim=-1)
self.E = torch.argmax(self.E, dim=-1)
self.X[node_mask == 0] = -1
self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = -1
else:
self.X = self.X * x_mask
self.E = self.E * e_mask1 * e_mask2
return self
def __repr__(self):
return (
f"X: {self.X.shape if type(self.X) == torch.Tensor else self.X} -- "
+ f"E: {self.E.shape if type(self.E) == torch.Tensor else self.E} -- "
+ f"y: {self.y.shape if type(self.y) == torch.Tensor else self.y}"
)
def split(self, node_mask):
"""Split a PlaceHolder representing a batch into a list of placeholders representing individual graphs."""
graph_list = []
batch_size = self.X.shape[0]
for i in range(batch_size):
n = torch.sum(node_mask[i], dim=0)
x = self.X[i, :n]
e = self.E[i, :n, :n]
y = self.y[i] if self.y is not None else None
graph_list.append(PlaceHolder(X=x, E=e, y=y))
return graph_list
|