File size: 2,967 Bytes
3f0529e |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import torch
def flatten_graph(node_embeddings, edge_embeddings, edge_index):
"""
Flattens the graph into a batch size one (with disconnected subgraphs for
each example) to be compatible with pytorch-geometric package.
Args:
node_embeddings: node embeddings in tuple form (scalar, vector)
- scalar: shape batch size x nodes x node_embed_dim
- vector: shape batch size x nodes x node_embed_dim x 3
edge_embeddings: edge embeddings of in tuple form (scalar, vector)
- scalar: shape batch size x edges x edge_embed_dim
- vector: shape batch size x edges x edge_embed_dim x 3
edge_index: shape batch_size x 2 (source node and target node) x edges
Returns:
node_embeddings: node embeddings in tuple form (scalar, vector)
- scalar: shape batch total_nodes x node_embed_dim
- vector: shape batch total_nodes x node_embed_dim x 3
edge_embeddings: edge embeddings of in tuple form (scalar, vector)
- scalar: shape batch total_edges x edge_embed_dim
- vector: shape batch total_edges x edge_embed_dim x 3
edge_index: shape 2 x total_edges
"""
x_s, x_v = node_embeddings
e_s, e_v = edge_embeddings
batch_size, N = x_s.shape[0], x_s.shape[1]
node_embeddings = (torch.flatten(x_s, 0, 1), torch.flatten(x_v, 0, 1))
edge_embeddings = (torch.flatten(e_s, 0, 1), torch.flatten(e_v, 0, 1))
edge_mask = torch.any(edge_index != -1, dim=1)
# Re-number the nodes by adding batch_idx * N to each batch
edge_index = edge_index + (torch.arange(batch_size, device=edge_index.device) *
N).unsqueeze(-1).unsqueeze(-1)
edge_index = edge_index.permute(1, 0, 2).flatten(1, 2)
edge_mask = edge_mask.flatten()
edge_index = edge_index[:, edge_mask]
edge_embeddings = (
edge_embeddings[0][edge_mask, :],
edge_embeddings[1][edge_mask, :]
)
return node_embeddings, edge_embeddings, edge_index
def unflatten_graph(node_embeddings, batch_size):
"""
Unflattens node embeddings.
Args:
node_embeddings: node embeddings in tuple form (scalar, vector)
- scalar: shape batch total_nodes x node_embed_dim
- vector: shape batch total_nodes x node_embed_dim x 3
batch_size: int
Returns:
node_embeddings: node embeddings in tuple form (scalar, vector)
- scalar: shape batch size x nodes x node_embed_dim
- vector: shape batch size x nodes x node_embed_dim x 3
"""
x_s, x_v = node_embeddings
x_s = x_s.reshape(batch_size, -1, x_s.shape[1])
x_v = x_v.reshape(batch_size, -1, x_v.shape[1], x_v.shape[2])
return (x_s, x_v)
|