Upsampling from global representation

[hkim716]

Hi Matt, I would like to ask about up-sampling from global pooled representations.

Let say I have encoder and decoder like this to reconstruct original input.
SplineConv->cluster=graclus->max_pool->global_max(or mean)_pool->bottle neck->up-sampled[cluster]->SplineConv.
When I have global_mean_pool, it takes average for each out_channel so we can get small dimensional bottle neck space like outputs from fully connected representations.
However, as you said, when I tried to use global_mean_pool, each node holds the same feature value, which makes it difficult to obtain meaningful up-sampled values during decoding.
If I must use global_mean_pool, (I need to have 1-dimensional representations at the bottle neck to train these representations with other dataset.), how would you up-sample from the global_mean_pool layer?

[rusty1s]

I think your adjacency matrix was fixed for each example, right? In this case this simplifies things a lot. You have two options:

1. Don't make use of global pooling and simply reshape your encoder output to view(batch_size, -1), and reshape in decoder to .view(num_coarsened_nodes * batch_size, -1) again.
2. Make use of global pooling but make the first transformation in the decoder have different parameters for each node it maps to:

lin = Linear(out_channels, num_coarsened_nodes * out_channels)
x = lin(z).view(num_coarsened_nodes * batch_size, -1)

[hkim716]

I tried to make batched_cluster. However, the shape of data.batch * max_clusters does not match with the size of torch.cat([cluster] * batch_size) for the second and third coarsened clusters. (In the case of depth=3)

def pre_process(train_dataset, depth):
    train_loader_pre = DataLoader(train_dataset, batch_size=1, shuffle=False)
    clusters = []
    data_batches = []
    for data in train_loader_pre:
        for i in range(depth):
            x, edge_index, edge_attr, pos = data.x.to(torch.float32), data.edge_index, data.edge_attr.to(torch.float32), data.pos.to(torch.float32)
            cluster = graclus(edge_index)
            clusters.append(cluster)
            data = max_pool(clusters[i], data, transform=transform)
        break
    return clusters

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

clusters = pre_process(train_dataset, 3)

batch_size = 5
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size)


for data in train_loader:
    data = data.to(device)
    batched_clusters = batched_cluster(data, batch_size, clusters, depth=3)
    break

batched_clusters[0] should work because torch.cat([clusters[0]] * batch_size) and data.batch have the same shape, but how could you generate batched_clusters[1], [2], etc when there's no information of the second coarsened data.batch?

[rusty1s]

Two things are missing here:

1. You need to run consecutive cluster after computing cluster.
2. You need to modify data.batch in each level of batched_cluster, e.g.:

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

[hkim716]

Thanks Matt, it worked, and I think batched_clusters are successfully generated.
However, there was another shape mismatch issue during forward function. (the mismatch did not happen when I tried with pre-process in a batch-wise fashion)
A part of my forward function looks like this.

def forward(self, data):
    x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr

    for i in range(len(self.down_path)):
        print(x.shape)
        x = self.down_path[i](x, edge_index, edge_attr)
        data.x = x
        print(self.batched_clusters[i].shape)
        data = max_pool(self.batched_clusters[i], data, transform=transform)
        print(data.x.shape)
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr

When I checked these three shapes, the results are

torch.Size([370, 32])
torch.Size([370])
torch.Size([206, 32])

torch.Size([206, 32])
torch.Size([210])

I think Size([206, 32]) should be Size([210, 32]), which is ([batched_clusters[1], hidden_channels]). I think 206 came from max_pool operation, which does not care about the next batched_cluster[1]. How could I make this shape same with the next batched_cluster?

The shapes of three batched_clusters generated here are

torch.Size([370])
torch.Size([210])
torch.Size([130])

and I found the difference of 210 and 206 was always 210-206=4, which was batch_size-1.
When I use batch_size=100, the difference was 99. It's weird..

[hkim716]

Oh, I found what the problem was!
batched_cluster = torch.cat([cluster] * batch_size).to(device) + batch * max_clusters
should have been batched_cluster = torch.cat([cluster] * batch_size).to(device) + batch * (max_clusters+1).

[rusty1s]

Oh, you are right. Thanks for fixing it :)