Batch cluster vectors for different graph data samples

[hkim716]

Hi Matt.

I'm trying to make a new graph autoencoder, where the inputs are graph dataset that has different number of nodes and edges. Some of graph data sample look like this:
Data(x=[5577, 1], edge_index=[2, 32816], y=[3], pos=[5577, 3], cluster=[6], edge_attr=[32816, 3])
Data(x=[5581, 1], edge_index=[2, 32836], y=[3], pos=[5581, 3], cluster=[6], edge_attr=[32836, 3])
...
So they may have different graph structures.

My graph autoencoder architecture looks like this :

class autoencoder(torch.nn.Module):
    def __init__(self, in_channel, hidden_channel, out_channel, fc1_hidden, x_lat,
                 depth, kernel_size, batch_size):
        super(autoencoder, self).__init__()
        
        self.in_channel = in_channel
        self.hidden_channel = hidden_channel
        self.out_channel = out_channel # = Latent_channel
        self.kernel_size = kernel_size
        self.batch_size = batch_size
        self.depth = depth
        self.fc1_hidden = fc1_hidden
        self.x_lat = x_lat
       
        self.encoder = Encoder(self.in_channel, self.hidden_channel, self.out_channel, self.fc1_hidden, self.x_lat,
                                                         self.depth, self.kernel_size, self.batch_size)
        
        self.decoder = Decoder(self.in_channel, self.hidden_channel, self.out_channel, self.fc1_hidden, self.x_lat,
                                                         self.depth, self.kernel_size, self.batch_size)
        
    def forward(self, data):

        b_cluster = batched_cluster(data, self.batch_size, data.cluster, self.depth)
        x, edge_indices, edge_attrs = self.encoder(data, b_cluster)
        
        emb = x
      
        x = self.decoder(x, edge_indices, edge_attrs, b_cluster)
        return emb, x, edge_indices, edge_attrs

Encoder and Decoder has six SplineConv and max_pool(clusters[i], data, transform=transform) layers, thus each graph sample stores cluster vectors, cluster=[6] in their data objects.

Current batched_cluster function was defined for mini-batch training. (You gave me an advice for this Thanks :))

def batched_cluster(data, batch_size, clusters, depth):
    batched_clusters = []
    batch = data.batch
    for i in range(depth):
        cluster = torch.tensor(clusters[0][i])
        max_clusters = int(cluster.max())
        batch = torch.repeat_interleave(torch.arange(batch_size), cluster.size(0))
        batched_cluster = torch.cat([cluster] * (batch_size)).to(device) + (batch * (max_clusters+1)).to(device)
        batched_clusters.append(batched_cluster)
    return batched_clusters

However, this batched_cluster worked for graph datasets that have identical graph structures i.e. edge_index, pos, edge_attr were the same.
In the case that we have arbitrary (or different) graph structures, how could I modify batched_cluster so I can use mini-batch training?

[rusty1s]

I think this is still an open problem. Most GNN-based autoencoder implementations I am aware of need to store graph connectivity for decoding, e.g., https://arxiv.org/abs/1807.10267. More recent approaches utilize different representations of meshes, e.g., https://arxiv.org/abs/1901.05103. It might be a good idea to do some literature research on any new methods in this regard.

[hkim716]

I already stored the node information for which nodes will be used for pooling.
For example, dataset1.cluster looks like this :

array([  0,   0,   1, ..., 468, 469, 456], dtype=int32), 
array([  0,   0,   1,   1,   1,   2,   2,   3,   3,   4,   5,   5,   6,
         7,   7,   8,   8,   9,  10,  11,  11,  12,  12,  13,  13,  14,
        14,  15,  15,  16,  16,  17,   0,  18,   2,   6,  19,   2,   6,
        ...
       128, 122, 137, 136, 130, 133, 133, 135, 124, 136, 126, 134, 132,
       135, 128, 122, 138, 138, 132, 137], dtype=int32), 
array([ 0,  1,  1,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  8,  9, 10, 11,
       11,  0,  0,  5, 10,  1,  4, 12,  2, 13,  2,  3,  4, 14, 14, 15, 14, ...    33, 13, 40, 13, 38, 27, 38, 39, 38], dtype=int32), 
array([ 0,  0,  1,  1,  2,  3,  4,  4,  4,  5,  6,  7,  0,  4,  0,  3,  8,
        1,  9, 10,  4,  9,  2,  3,  9,  8, 10,  3,  8,  9,  9,  9, 11,  8,
       10,  3, 10,  2, 11, 11, 10], dtype=int32), array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)]

What I'm trying to do is I want to train the model with batch_size>1. Since each graph sample has different pooling nodes, I think I need to modify my code.
When I use batch_size=2, I have error code like this:

---> 26         x, edge_indices, edge_attrs = self.encoder(data, b_cluster)
     27 
     28         emb = x

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/scratch/4567855.sched-torque.pace.gatech.edu/ipykernel_190553/896632367.py in forward(self, data, clusters)
     52 
---> 53             data = max_pool(clusters[i], data, transform=transform)
     55             x, edge_index, edge_attr = data.x.to(torch.float32), data.edge_index, data.edge_attr.to(torch.float32)

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch_geometric/nn/pool/max_pool.py in max_pool(cluster, data, transform)
     67     cluster, perm = consecutive_cluster(cluster)
     68 
---> 69     x = None if data.x is None else _max_pool_x(cluster, data.x)
     70     index, attr = pool_edge(cluster, data.edge_index, data.edge_attr)
     71     batch = None if data.batch is None else pool_batch(perm, data.batch)

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch_geometric/nn/pool/max_pool.py in _max_pool_x(cluster, x, size)
     11 
     12 def _max_pool_x(cluster, x, size: Optional[int] = None):
---> 13     return scatter(x, cluster, dim=0, dim_size=size, reduce='max')
     14 
     15 

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch_scatter/scatter.py in scatter(src, index, dim, out, dim_size, reduce)
    159         return scatter_min(src, index, dim, out, dim_size)[0]
    160     elif reduce == 'max':
--> 161         return scatter_max(src, index, dim, out, dim_size)[0]
    162     else:
    163         raise ValueError

~/.conda/envs/gnb_df_pyg/lib/python3.7/site-packages/torch_scatter/scatter.py in scatter_max(src, index, dim, out, dim_size)
     71         out: Optional[torch.Tensor] = None,
     72         dim_size: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
---> 73     return torch.ops.torch_scatter.scatter_max(src, index, dim, out, dim_size)
     74 
     75 

RuntimeError: The expanded size of the tensor (11157) must match the existing size (11160) at non-singleton dimension 0.  Target sizes: [11157, 4].  Tensor sizes: [11160, 1]

[rusty1s]

I think you will need to modify batched_clusters such that it increases cluster indices correctly for a varying number of clusters across graphs:

def batched_cluster(data, batch_size, clusters, depth):
    batched_clusters = []
    batch = data.batch
    for i in range(depth):
        i = 0
        all_clusters = []
        for cluster in [c[i] for c in clusters]:
            all_clusters.append(cluster + i)
            i += int(cluster.max()) + 1
        cluster = torch.cat(all_clusters)